FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Add support for disable_builtin_metrics in SpannerConnection and connect · google/python-spanner-orm@7699d2e · GitHub

Commit 7699d2e

Browse files
committed
Add support for disable_builtin_metrics in SpannerConnection and connect
Allow passing disable_builtin_metrics to SpannerConnection and connect helper functions to be forwarded to google.cloud.spanner.Client.
1 parent b9b9191 commit 7699d2e

3 files changed

Lines changed: 68 additions & 13 deletions

File tree

‎spanner_orm/admin/api.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,15 @@
1313
# limitations under the License.
1414
"""Class that handles API calls to Spanner that deal with table metadata."""
1515

16-
from typing import Iterable, Optional
16+
from typing import Any, Dict, Iterable, Optional, Union
1717
import warnings
1818

19-
from spanner_orm import api
20-
from spanner_orm import error
21-
19+
from google.api_core import client_options as api_client_options
2220
from google.auth import credentials as auth_credentials
2321
from google.cloud.spanner_v1 import database as spanner_database
2422
from google.cloud.spanner_v1 import pool as spanner_pool
23+
from spanner_orm import api
24+
from spanner_orm import error
2525

2626

2727
class SpannerAdminApi(api.SpannerReadApi, api.SpannerWriteApi):
@@ -58,7 +58,11 @@ def connect(instance: str,
5858
project: Optional[str] = None,
5959
credentials: Optional[auth_credentials.Credentials] = None,
6060
pool: Optional[spanner_pool.AbstractSessionPool] = None,
61-
create_ddl: Optional[Iterable[str]] = None) -> SpannerAdminApi:
61+
create_ddl: Optional[Iterable[str]] = None,
62+
*,
63+
client_options: Union[api_client_options.ClientOptions,
64+
Dict[Any, Any], None] = None,
65+
disable_builtin_metrics: Optional[bool] = None) -> SpannerAdminApi:
6266
"""Connects the global Spanner admin API to a Spanner database.
6367
6468
Deprecated in favor of from_connection().
@@ -72,7 +76,9 @@ def connect(instance: str,
7276
project=project,
7377
credentials=credentials,
7478
pool=pool,
75-
create_ddl=create_ddl)
79+
create_ddl=create_ddl,
80+
client_options=client_options,
81+
disable_builtin_metrics=disable_builtin_metrics)
7682
return from_connection(connection)
7783

7884

‎spanner_orm/api.py‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ def __init__(
117117
*,
118118
client_options: Union[api_client_options.ClientOptions, Dict[Any, Any],
119119
None] = None,
120+
disable_builtin_metrics: Optional[bool] = None,
120121
):
121122
"""Connects to the specified Spanner database."""
122123
self._instance = instance
@@ -126,15 +127,20 @@ def __init__(
126127
self._pool = pool
127128
self._create_ddl = create_ddl
128129
self._client_options = client_options
130+
self._disable_builtin_metrics = disable_builtin_metrics
129131
self.connect()
130132

131133
def connect(self):
132134
"""Establish a new connection to the specified Spanner database."""
133-
client = spanner.Client(
134-
project=self._project,
135-
credentials=self._credentials,
136-
client_options=self._client_options,
137-
)
135+
client_kwargs = {
136+
'project': self._project,
137+
'credentials': self._credentials,
138+
'client_options': self._client_options,
139+
}
140+
if self._disable_builtin_metrics is not None:
141+
client_kwargs['disable_builtin_metrics'] = self._disable_builtin_metrics
142+
143+
client = spanner.Client(**client_kwargs)
138144
instance = client.instance(self._instance)
139145
self.database = instance.database(
140146
self._database, pool=self._pool, ddl_statements=self._create_ddl or ())
@@ -164,7 +170,11 @@ def connect(
164170
database: str,
165171
project: Optional[str] = None,
166172
credentials: Optional[auth_credentials.Credentials] = None,
167-
pool: Optional[spanner_pool.AbstractSessionPool] = None) -> SpannerApi:
173+
pool: Optional[spanner_pool.AbstractSessionPool] = None,
174+
*,
175+
client_options: Union[api_client_options.ClientOptions, Dict[Any, Any],
176+
None] = None,
177+
disable_builtin_metrics: Optional[bool] = None) -> SpannerApi:
168178
"""Connects to the Spanner database and sets the global spanner_api.
169179
170180
Deprecated in favor of from_connection().
@@ -174,7 +184,13 @@ def connect(
174184
'Please use '
175185
'spanner_orm.from_connection(spanner_orm.SpannerConnection(...))'))
176186
connection = SpannerConnection(
177-
instance, database, project=project, credentials=credentials, pool=pool)
187+
instance,
188+
database,
189+
project=project,
190+
credentials=credentials,
191+
pool=pool,
192+
client_options=client_options,
193+
disable_builtin_metrics=disable_builtin_metrics)
178194
return from_connection(connection)
179195

180196

‎spanner_orm/tests/api_test.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,39 @@ def test_connection_args(self, client):
7373
client.mock_calls,
7474
)
7575

76+
@mock.patch.object(spanner, 'Client', autospec=True, spec_set=True)
77+
def test_connection_args_with_disable_builtin_metrics(self, client):
78+
client.return_value.instance.return_value.database.return_value = (
79+
'fake-database')
80+
connection = api.SpannerConnection(
81+
instance='some-instance',
82+
database='some-database',
83+
project='some-project',
84+
credentials='fake-credentials',
85+
pool='fake-pool',
86+
create_ddl=('fake-ddl',),
87+
client_options=dict(fake='options'),
88+
disable_builtin_metrics=True,
89+
)
90+
self.assertEqual('fake-database', connection.database)
91+
self.assertSequenceEqual(
92+
(
93+
mock.call(
94+
project='some-project',
95+
credentials='fake-credentials',
96+
client_options=dict(fake='options'),
97+
disable_builtin_metrics=True,
98+
),
99+
mock.call().instance('some-instance'),
100+
mock.call().instance().database(
101+
'some-database',
102+
pool='fake-pool',
103+
ddl_statements=('fake-ddl',),
104+
),
105+
),
106+
client.mock_calls,
107+
)
108+
76109
@mock.patch('google.cloud.spanner.Client')
77110
def test_api_connection(self, client):
78111
connection = self.mock_connection(client)

0 commit comments

Comments
 (0)

Back | FazBrowse Home | New Git URL