Skip to content

Commit bbc6b36

Browse files
authored
feat: support quota project override via client options (#496)
* feat: support quota project override via client options * chore(deps): bump api-core version * Update setup.py.j2
1 parent f86a47b commit bbc6b36

File tree

7 files changed

+78
-20
lines changed

7 files changed

+78
-20
lines changed

gapic/templates/%namespace/%name_%version/%sub/services/%service/client.py.j2

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
214214
scopes=client_options.scopes,
215215
api_mtls_endpoint=client_options.api_endpoint,
216216
client_cert_source=client_options.client_cert_source,
217+
quota_project_id=client_options.quota_project_id,
217218
)
218219

219220
{% for method in service.methods.values() -%}

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/base.py.j2

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ class {{ service.name }}Transport(abc.ABC):
3333
credentials: credentials.Credentials = None,
3434
credentials_file: typing.Optional[str] = None,
3535
scopes: typing.Optional[typing.Sequence[str]] = AUTH_SCOPES,
36+
quota_project_id: typing.Optional[str] = None,
3637
**kwargs,
3738
) -> None:
3839
"""Instantiate the transport.
@@ -49,6 +50,8 @@ class {{ service.name }}Transport(abc.ABC):
4950
be loaded with :func:`google.auth.load_credentials_from_file`.
5051
This argument is mutually exclusive with credentials.
5152
scope (Optional[Sequence[str]]): A list of scopes.
53+
quota_project_id (Optional[str]): An optional project to use for billing
54+
and quota.
5255
"""
5356
# Save the hostname. Default to port 443 (HTTPS) if none is specified.
5457
if ':' not in host:
@@ -61,9 +64,14 @@ class {{ service.name }}Transport(abc.ABC):
6164
raise exceptions.DuplicateCredentialArgs("'credentials_file' and 'credentials' are mutually exclusive")
6265

6366
if credentials_file is not None:
64-
credentials, _ = auth.load_credentials_from_file(credentials_file, scopes=scopes)
67+
credentials, _ = auth.load_credentials_from_file(
68+
credentials_file,
69+
scopes=scopes,
70+
quota_project_id=quota_project_id
71+
)
72+
6573
elif credentials is None:
66-
credentials, _ = auth.default(scopes=scopes)
74+
credentials, _ = auth.default(scopes=scopes, quota_project_id=quota_project_id)
6775

6876
# Save the credentials.
6977
self._credentials = credentials

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/grpc.py.j2

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
4444
scopes: Sequence[str] = None,
4545
channel: grpc.Channel = None,
4646
api_mtls_endpoint: str = None,
47-
client_cert_source: Callable[[], Tuple[bytes, bytes]] = None) -> None:
47+
client_cert_source: Callable[[], Tuple[bytes, bytes]] = None,
48+
quota_project_id: Optional[str] = None) -> None:
4849
"""Instantiate the transport.
4950

5051
Args:
@@ -71,6 +72,8 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
7172
callback to provide client SSL certificate bytes and private key
7273
bytes, both in PEM format. It is ignored if ``api_mtls_endpoint``
7374
is None.
75+
quota_project_id (Optional[str]): An optional project to use for billing
76+
and quota.
7477

7578
Raises:
7679
google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport
@@ -89,7 +92,7 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
8992
host = api_mtls_endpoint if ":" in api_mtls_endpoint else api_mtls_endpoint + ":443"
9093

9194
if credentials is None:
92-
credentials, _ = auth.default(scopes=self.AUTH_SCOPES)
95+
credentials, _ = auth.default(scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id)
9396

9497
# Create SSL credentials with client_cert_source or application
9598
# default SSL credentials.
@@ -108,14 +111,16 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
108111
credentials_file=credentials_file,
109112
ssl_credentials=ssl_credentials,
110113
scopes=scopes or self.AUTH_SCOPES,
114+
quota_project_id=quota_project_id,
111115
)
112116

113117
# Run the base constructor.
114118
super().__init__(
115119
host=host,
116120
credentials=credentials,
117121
credentials_file=credentials_file,
118-
scopes=scopes or self.AUTH_SCOPES
122+
scopes=scopes or self.AUTH_SCOPES,
123+
quota_project_id=quota_project_id,
119124
)
120125

121126
self._stubs = {} # type: Dict[str, Callable]
@@ -126,6 +131,7 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
126131
credentials: credentials.Credentials = None,
127132
credentials_file: str = None,
128133
scopes: Optional[Sequence[str]] = None,
134+
quota_project_id: Optional[str] = None,
129135
**kwargs) -> grpc.Channel:
130136
"""Create and return a gRPC channel object.
131137
Args:
@@ -141,6 +147,8 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
141147
scopes (Optional[Sequence[str]]): A optional list of scopes needed for this
142148
service. These are only used when credentials are not specified and
143149
are passed to :func:`google.auth.default`.
150+
quota_project_id (Optional[str]): An optional project to use for billing
151+
and quota.
144152
kwargs (Optional[dict]): Keyword arguments, which are passed to the
145153
channel creation.
146154
Returns:
@@ -156,6 +164,7 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
156164
credentials=credentials,
157165
credentials_file=credentials_file,
158166
scopes=scopes,
167+
quota_project_id=quota_project_id,
159168
**kwargs
160169
)
161170

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/grpc_asyncio.py.j2

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
4545
credentials: credentials.Credentials = None,
4646
credentials_file: Optional[str] = None,
4747
scopes: Optional[Sequence[str]] = None,
48+
quota_project_id: Optional[str] = None,
4849
**kwargs) -> aio.Channel:
4950
"""Create and return a gRPC AsyncIO channel object.
5051
Args:
@@ -60,6 +61,8 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
6061
scopes (Optional[Sequence[str]]): A optional list of scopes needed for this
6162
service. These are only used when credentials are not specified and
6263
are passed to :func:`google.auth.default`.
64+
quota_project_id (Optional[str]): An optional project to use for billing
65+
and quota.
6366
kwargs (Optional[dict]): Keyword arguments, which are passed to the
6467
channel creation.
6568
Returns:
@@ -71,6 +74,7 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
7174
credentials=credentials,
7275
credentials_file=credentials_file,
7376
scopes=scopes,
77+
quota_project_id=quota_project_id,
7478
**kwargs
7579
)
7680

@@ -81,7 +85,9 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
8185
scopes: Optional[Sequence[str]] = None,
8286
channel: aio.Channel = None,
8387
api_mtls_endpoint: str = None,
84-
client_cert_source: Callable[[], Tuple[bytes, bytes]] = None) -> None:
88+
client_cert_source: Callable[[], Tuple[bytes, bytes]] = None,
89+
quota_project_id=None,
90+
) -> None:
8591
"""Instantiate the transport.
8692

8793
Args:
@@ -109,6 +115,8 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
109115
callback to provide client SSL certificate bytes and private key
110116
bytes, both in PEM format. It is ignored if ``api_mtls_endpoint``
111117
is None.
118+
quota_project_id (Optional[str]): An optional project to use for billing
119+
and quota.
112120

113121
Raises:
114122
google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport
@@ -143,14 +151,16 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
143151
credentials_file=credentials_file,
144152
ssl_credentials=ssl_credentials,
145153
scopes=scopes or self.AUTH_SCOPES,
154+
quota_project_id=quota_project_id,
146155
)
147156

148157
# Run the base constructor.
149158
super().__init__(
150159
host=host,
151160
credentials=credentials,
152161
credentials_file=credentials_file,
153-
scopes=scopes or self.AUTH_SCOPES
162+
scopes=scopes or self.AUTH_SCOPES,
163+
quota_project_id=quota_project_id,
154164
)
155165

156166
self._stubs = {}

gapic/templates/setup.py.j2

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ setuptools.setup(
1616
platforms='Posix; MacOS X; Windows',
1717
include_package_data=True,
1818
install_requires=(
19-
'google-api-core[grpc] >= 1.21.0, < 2.0.0dev',
19+
'google-api-core[grpc] >= 1.22.0, < 2.0.0dev',
2020
'libcst >= 0.2.5',
2121
'proto-plus >= 1.1.0',
2222
{%- if api.requires_package(('google', 'iam', 'v1')) %}

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
112112
scopes=None,
113113
api_mtls_endpoint="squid.clam.whelk",
114114
client_cert_source=None,
115+
quota_project_id=None,
115116
)
116117

117118
# Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS is
@@ -127,6 +128,7 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
127128
scopes=None,
128129
api_mtls_endpoint=client.DEFAULT_ENDPOINT,
129130
client_cert_source=None,
131+
quota_project_id=None,
130132
)
131133

132134
# Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS is
@@ -142,6 +144,7 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
142144
scopes=None,
143145
api_mtls_endpoint=client.DEFAULT_MTLS_ENDPOINT,
144146
client_cert_source=None,
147+
quota_project_id=None,
145148
)
146149

147150
# Check the case api_endpoint is not provided, GOOGLE_API_USE_MTLS is
@@ -158,6 +161,7 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
158161
scopes=None,
159162
api_mtls_endpoint=client.DEFAULT_MTLS_ENDPOINT,
160163
client_cert_source=client_cert_source_callback,
164+
quota_project_id=None,
161165

162166
)
163167

@@ -175,6 +179,7 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
175179
scopes=None,
176180
api_mtls_endpoint=client.DEFAULT_MTLS_ENDPOINT,
177181
client_cert_source=None,
182+
quota_project_id=None,
178183
)
179184

180185
# Check the case api_endpoint is not provided, GOOGLE_API_USE_MTLS is
@@ -191,15 +196,29 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
191196
scopes=None,
192197
api_mtls_endpoint=client.DEFAULT_ENDPOINT,
193198
client_cert_source=None,
199+
quota_project_id=None,
194200
)
195201

196202
# Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS has
197203
# unsupported value.
198-
os.environ["GOOGLE_API_USE_MTLS"] = "Unsupported"
199-
with pytest.raises(MutualTLSChannelError):
200-
client = client_class()
204+
with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS": "Unsupported"}):
205+
with pytest.raises(MutualTLSChannelError):
206+
client = client_class()
201207

202-
del os.environ["GOOGLE_API_USE_MTLS"]
208+
# Check the case quota_project_id is provided
209+
options = client_options.ClientOptions(quota_project_id="octopus")
210+
with mock.patch.object(transport_class, '__init__') as patched:
211+
patched.return_value = None
212+
client = client_class(client_options=options)
213+
patched.assert_called_once_with(
214+
credentials=None,
215+
credentials_file=None,
216+
host=client.DEFAULT_ENDPOINT,
217+
scopes=None,
218+
api_mtls_endpoint=client.DEFAULT_ENDPOINT,
219+
client_cert_source=None,
220+
quota_project_id="octopus",
221+
)
203222

204223

205224
@pytest.mark.parametrize("client_class,transport_class,transport_name", [
@@ -221,6 +240,7 @@ def test_{{ service.client_name|snake_case }}_client_options_scopes(client_class
221240
scopes=["1", "2"],
222241
api_mtls_endpoint=client.DEFAULT_ENDPOINT,
223242
client_cert_source=None,
243+
quota_project_id=None,
224244
)
225245

226246

@@ -243,6 +263,7 @@ def test_{{ service.client_name|snake_case }}_client_options_credentials_file(cl
243263
scopes=None,
244264
api_mtls_endpoint=client.DEFAULT_ENDPOINT,
245265
client_cert_source=None,
266+
quota_project_id=None,
246267
)
247268

248269

@@ -259,6 +280,7 @@ def test_{{ service.client_name|snake_case }}_client_options_from_dict():
259280
scopes=None,
260281
api_mtls_endpoint="squid.clam.whelk",
261282
client_cert_source=None,
283+
quota_project_id=None,
262284
)
263285

264286

@@ -1001,12 +1023,15 @@ def test_{{ service.name|snake_case }}_base_transport_with_credentials_file():
10011023
load_creds.return_value = (credentials.AnonymousCredentials(), None)
10021024
transport = transports.{{ service.name }}Transport(
10031025
credentials_file="credentials.json",
1026+
quota_project_id="octopus",
10041027
)
10051028
load_creds.assert_called_once_with("credentials.json", scopes=(
10061029
{%- for scope in service.oauth_scopes %}
10071030
'{{ scope }}',
10081031
{%- endfor %}
1009-
))
1032+
),
1033+
quota_project_id="octopus",
1034+
)
10101035

10111036

10121037
def test_{{ service.name|snake_case }}_auth_adc():
@@ -1017,22 +1042,23 @@ def test_{{ service.name|snake_case }}_auth_adc():
10171042
adc.assert_called_once_with(scopes=(
10181043
{%- for scope in service.oauth_scopes %}
10191044
'{{ scope }}',
1020-
{%- endfor %}
1021-
))
1045+
{%- endfor %}),
1046+
quota_project_id=None,
1047+
)
10221048

10231049

10241050
def test_{{ service.name|snake_case }}_transport_auth_adc():
10251051
# If credentials and host are not provided, the transport class should use
10261052
# ADC credentials.
10271053
with mock.patch.object(auth, 'default') as adc:
10281054
adc.return_value = (credentials.AnonymousCredentials(), None)
1029-
transports.{{ service.name }}GrpcTransport(host="squid.clam.whelk")
1055+
transports.{{ service.name }}GrpcTransport(host="squid.clam.whelk", quota_project_id="octopus")
10301056
adc.assert_called_once_with(scopes=(
10311057
{%- for scope in service.oauth_scopes %}
10321058
'{{ scope }}',
1033-
{%- endfor %}
1034-
))
1035-
1059+
{%- endfor %}),
1060+
quota_project_id="octopus",
1061+
)
10361062

10371063
def test_{{ service.name|snake_case }}_host_no_port():
10381064
{% with host = (service.host|default('localhost', true)).split(':')[0] -%}
@@ -1122,6 +1148,7 @@ def test_{{ service.name|snake_case }}_grpc_transport_channel_mtls_with_client_c
11221148
{%- endfor %}
11231149
),
11241150
ssl_credentials=mock_ssl_cred,
1151+
quota_project_id=None,
11251152
)
11261153
assert transport.grpc_channel == mock_grpc_channel
11271154

@@ -1160,6 +1187,7 @@ def test_{{ service.name|snake_case }}_grpc_asyncio_transport_channel_mtls_with_
11601187
{%- endfor %}
11611188
),
11621189
ssl_credentials=mock_ssl_cred,
1190+
quota_project_id=None,
11631191
)
11641192
assert transport.grpc_channel == mock_grpc_channel
11651193

@@ -1200,6 +1228,7 @@ def test_{{ service.name|snake_case }}_grpc_transport_channel_mtls_with_adc(
12001228
{%- endfor %}
12011229
),
12021230
ssl_credentials=mock_ssl_cred,
1231+
quota_project_id=None,
12031232
)
12041233
assert transport.grpc_channel == mock_grpc_channel
12051234

@@ -1240,6 +1269,7 @@ def test_{{ service.name|snake_case }}_grpc_asyncio_transport_channel_mtls_with_
12401269
{%- endfor %}
12411270
),
12421271
ssl_credentials=mock_ssl_cred,
1272+
quota_project_id=None,
12431273
)
12441274
assert transport.grpc_channel == mock_grpc_channel
12451275

requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
click==7.1.2
2-
google-api-core==1.21.0
2+
google-api-core==1.22.0
33
googleapis-common-protos==1.52.0
44
jinja2==2.11.2
55
MarkupSafe==1.1.1

0 commit comments

Comments
 (0)