|
32 | 32 | import grpc |
33 | 33 |
|
34 | 34 | from google.api_core.gapic_v1 import client_info |
| 35 | +import google.auth |
35 | 36 |
|
36 | 37 | from google.cloud import bigtable_v2 |
37 | 38 | from google.cloud import bigtable_admin_v2 |
|
69 | 70 |
|
70 | 71 | def _create_gapic_client(client_class, client_options=None, transport=None): |
71 | 72 | def inner(self): |
72 | | - if self._emulator_host is None: |
73 | | - return client_class( |
74 | | - credentials=None, |
75 | | - client_info=self._client_info, |
76 | | - client_options=client_options, |
77 | | - transport=transport, |
78 | | - ) |
79 | | - else: |
80 | | - return client_class( |
81 | | - channel=self._emulator_channel, client_info=self._client_info |
82 | | - ) |
| 73 | + return client_class( |
| 74 | + credentials=None, |
| 75 | + client_info=self._client_info, |
| 76 | + client_options=client_options, |
| 77 | + transport=transport, |
| 78 | + ) |
83 | 79 |
|
84 | 80 | return inner |
85 | 81 |
|
@@ -166,16 +162,6 @@ def __init__( |
166 | 162 | self._admin = bool(admin) |
167 | 163 | self._client_info = client_info |
168 | 164 | self._emulator_host = os.getenv(BIGTABLE_EMULATOR) |
169 | | - self._emulator_channel = None |
170 | | - |
171 | | - if self._emulator_host is not None: |
172 | | - self._emulator_channel = grpc.insecure_channel( |
173 | | - target=self._emulator_host, |
174 | | - options={ |
175 | | - "grpc.keepalive_time_ms": 30000, |
176 | | - "grpc.keepalive_timeout_ms": 10000, |
177 | | - }.items(), |
178 | | - ) |
179 | 165 |
|
180 | 166 | if channel is not None: |
181 | 167 | warnings.warn( |
@@ -208,22 +194,76 @@ def _get_scopes(self): |
208 | 194 |
|
209 | 195 | return scopes |
210 | 196 |
|
| 197 | + def _emulator_channel(self, transport, options): |
| 198 | + """ |
| 199 | + Creates a channel using self._credentials in a similar way to grpc.secure_channel but |
| 200 | + using grpc.local_channel_credentials() rather than grpc.ssh_channel_credentials() |
| 201 | + to allow easy connection to a local emulator. |
| 202 | + :return: grpc.Channel or grpc.aio.Channel |
| 203 | + """ |
| 204 | + # TODO: Implement a special credentials type for emulator and use |
| 205 | + # "transport.create_channel" to create gRPC channels once google-auth |
| 206 | + # extends it's allowed credentials types. |
| 207 | + # Note: this code also exists in the firestore client. |
| 208 | + if "GrpcAsyncIOTransport" in str(transport.__name__): |
| 209 | + return grpc.aio.secure_channel( |
| 210 | + self._emulator_host, |
| 211 | + self._local_composite_credentials(), |
| 212 | + options=options, |
| 213 | + ) |
| 214 | + else: |
| 215 | + return grpc.secure_channel( |
| 216 | + self._emulator_host, |
| 217 | + self._local_composite_credentials(), |
| 218 | + options=options, |
| 219 | + ) |
| 220 | + |
| 221 | + def _local_composite_credentials(self): |
| 222 | + """ |
| 223 | + Creates the credentials for the local emulator channel |
| 224 | + :return: grpc.ChannelCredentials |
| 225 | + """ |
| 226 | + credentials = google.auth.credentials.with_scopes_if_required( |
| 227 | + self._credentials, None |
| 228 | + ) |
| 229 | + request = google.auth.transport.requests.Request() |
| 230 | + |
| 231 | + # Create the metadata plugin for inserting the authorization header. |
| 232 | + metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin( |
| 233 | + credentials, request |
| 234 | + ) |
| 235 | + |
| 236 | + # Create a set of grpc.CallCredentials using the metadata plugin. |
| 237 | + google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin) |
| 238 | + |
| 239 | + # Using the local_credentials to allow connection to emulator |
| 240 | + local_credentials = grpc.local_channel_credentials() |
| 241 | + |
| 242 | + # Combine the local credentials and the authorization credentials. |
| 243 | + return grpc.composite_channel_credentials( |
| 244 | + local_credentials, google_auth_credentials |
| 245 | + ) |
| 246 | + |
211 | 247 | def _create_gapic_client_channel(self, client_class, grpc_transport): |
| 248 | + options = { |
| 249 | + "grpc.max_send_message_length": -1, |
| 250 | + "grpc.max_receive_message_length": -1, |
| 251 | + "grpc.keepalive_time_ms": 30000, |
| 252 | + "grpc.keepalive_timeout_ms": 10000, |
| 253 | + }.items() |
212 | 254 | if self._client_options and self._client_options.api_endpoint: |
213 | 255 | api_endpoint = self._client_options.api_endpoint |
214 | 256 | else: |
215 | 257 | api_endpoint = client_class.DEFAULT_ENDPOINT |
216 | 258 |
|
217 | | - channel = grpc_transport.create_channel( |
218 | | - host=api_endpoint, |
219 | | - credentials=self._credentials, |
220 | | - options={ |
221 | | - "grpc.max_send_message_length": -1, |
222 | | - "grpc.max_receive_message_length": -1, |
223 | | - "grpc.keepalive_time_ms": 30000, |
224 | | - "grpc.keepalive_timeout_ms": 10000, |
225 | | - }.items(), |
226 | | - ) |
| 259 | + channel = None |
| 260 | + if self._emulator_host is not None: |
| 261 | + api_endpoint = self._emulator_host |
| 262 | + channel = self._emulator_channel(grpc_transport, options) |
| 263 | + else: |
| 264 | + channel = grpc_transport.create_channel( |
| 265 | + host=api_endpoint, credentials=self._credentials, options=options, |
| 266 | + ) |
227 | 267 | transport = grpc_transport(channel=channel, host=api_endpoint) |
228 | 268 | return transport |
229 | 269 |
|
|
0 commit comments