Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/google/cloud/bigquery/dbapi/connection.py: 29%

1# Copyright 2017 Google LLC 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Connection for the Google BigQuery DB-API.""" 

16 

17import weakref 

18 

19from google.cloud import bigquery 

20from google.cloud.bigquery.dbapi import cursor 

21from google.cloud.bigquery.dbapi import _helpers 

22 

23 

24@_helpers.raise_on_closed("Operating on a closed connection.") 

25class Connection(object): 

26 """DB-API Connection to Google BigQuery. 

27 

28 Args: 

29 client (Optional[google.cloud.bigquery.Client]): 

30 A REST API client used to connect to BigQuery. If not passed, a 

31 client is created using default options inferred from the environment. 

32 bqstorage_client(\ 

33 Optional[google.cloud.bigquery_storage_v1.BigQueryReadClient] \ 

34 ): 

35 A client that uses the faster BigQuery Storage API to fetch rows from 

36 BigQuery. If not passed, it is created using the same credentials 

37 as ``client`` (provided that BigQuery Storage dependencies are installed). 

38 prefer_bqstorage_client (Optional[bool]): 

39 Prefer the BigQuery Storage client over the REST client. If Storage 

40 client isn't available, fall back to the REST client. Defaults to 

41 ``True``. 

42 """ 

43 

44 def __init__( 

45 self, 

46 client=None, 

47 bqstorage_client=None, 

48 prefer_bqstorage_client=True, 

49 ): 

50 if client is None: 

51 client = bigquery.Client() 

52 self._owns_client = True 

53 else: 

54 self._owns_client = False 

55 

56 # A warning is already raised by the BQ Storage client factory factory if 

57 # instantiation fails, or if the given BQ Storage client instance is outdated. 

58 if not prefer_bqstorage_client: 

59 bqstorage_client = None 

60 self._owns_bqstorage_client = False 

61 elif bqstorage_client is None: 

62 bqstorage_client = client._ensure_bqstorage_client() 

63 self._owns_bqstorage_client = bqstorage_client is not None 

64 else: 

65 self._owns_bqstorage_client = False 

66 bqstorage_client = client._ensure_bqstorage_client(bqstorage_client) 

67 

68 self._client = client 

69 self._bqstorage_client = bqstorage_client 

70 

71 self._closed = False 

72 self._cursors_created = weakref.WeakSet() 

73 

74 def close(self): 

75 """Close the connection and any cursors created from it. 

76 

77 Any BigQuery clients explicitly passed to the constructor are *not* 

78 closed, only those created by the connection instance itself. 

79 """ 

80 self._closed = True 

81 

82 if self._owns_client: 

83 self._client.close() 

84 

85 if self._owns_bqstorage_client: 

86 # There is no close() on the BQ Storage client itself. 

87 self._bqstorage_client._transport.grpc_channel.close() 

88 

89 for cursor_ in self._cursors_created: 

90 if not cursor_._closed: 

91 cursor_.close() 

92 

93 def commit(self): 

94 """No-op, but for consistency raise an error if connection is closed.""" 

95 

96 def cursor(self): 

97 """Return a new cursor object. 

98 

99 Returns: 

100 google.cloud.bigquery.dbapi.Cursor: A DB-API cursor that uses this connection. 

101 """ 

102 new_cursor = cursor.Cursor(self) 

103 self._cursors_created.add(new_cursor) 

104 return new_cursor 

105 

106 

107def connect(client=None, bqstorage_client=None, prefer_bqstorage_client=True): 

108 """Construct a DB-API connection to Google BigQuery. 

109 

110 Args: 

111 client (Optional[google.cloud.bigquery.Client]): 

112 A REST API client used to connect to BigQuery. If not passed, a 

113 client is created using default options inferred from the environment. 

114 bqstorage_client(\ 

115 Optional[google.cloud.bigquery_storage_v1.BigQueryReadClient] \ 

116 ): 

117 A client that uses the faster BigQuery Storage API to fetch rows from 

118 BigQuery. If not passed, it is created using the same credentials 

119 as ``client`` (provided that BigQuery Storage dependencies are installed). 

120 prefer_bqstorage_client (Optional[bool]): 

121 Prefer the BigQuery Storage client over the REST client. If Storage 

122 client isn't available, fall back to the REST client. Defaults to 

123 ``True``. 

124 

125 Returns: 

126 google.cloud.bigquery.dbapi.Connection: A new DB-API connection to BigQuery. 

127 """ 

128 return Connection(client, bqstorage_client, prefer_bqstorage_client)