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

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 

39 If both clients are available, ``bqstorage_client`` is used for 

40 fetching query results. 

41 """ 

42 

43 def __init__(self, client=None, bqstorage_client=None): 

44 if client is None: 

45 client = bigquery.Client() 

46 self._owns_client = True 

47 else: 

48 self._owns_client = False 

49 

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

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

52 if bqstorage_client is None: 

53 bqstorage_client = client._ensure_bqstorage_client() 

54 self._owns_bqstorage_client = bqstorage_client is not None 

55 else: 

56 self._owns_bqstorage_client = False 

57 bqstorage_client = client._ensure_bqstorage_client(bqstorage_client) 

58 

59 self._client = client 

60 self._bqstorage_client = bqstorage_client 

61 

62 self._closed = False 

63 self._cursors_created = weakref.WeakSet() 

64 

65 def close(self): 

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

67 

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

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

70 """ 

71 self._closed = True 

72 

73 if self._owns_client: 

74 self._client.close() 

75 

76 if self._owns_bqstorage_client: 

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

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

79 

80 for cursor_ in self._cursors_created: 

81 if not cursor_._closed: 

82 cursor_.close() 

83 

84 def commit(self): 

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

86 

87 def cursor(self): 

88 """Return a new cursor object. 

89 

90 Returns: 

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

92 """ 

93 new_cursor = cursor.Cursor(self) 

94 self._cursors_created.add(new_cursor) 

95 return new_cursor 

96 

97 

98def connect(client=None, bqstorage_client=None): 

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

100 

101 Args: 

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

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

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

105 bqstorage_client(\ 

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

107 ): 

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

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

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

111 

112 If both clients are available, ``bqstorage_client`` is used for 

113 fetching query results. 

114 

115 Returns: 

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

117 """ 

118 return Connection(client, bqstorage_client)