Skip to content
This repository was archived by the owner on Mar 2, 2026. It is now read-only.
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions google/cloud/datastore_admin_v1/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

from .types.datastore_admin import CommonMetadata
from .types.datastore_admin import CreateIndexRequest
from .types.datastore_admin import DatastoreFirestoreMigrationMetadata
from .types.datastore_admin import DeleteIndexRequest
from .types.datastore_admin import EntityFilter
from .types.datastore_admin import ExportEntitiesMetadata
Expand All @@ -33,12 +34,17 @@
from .types.datastore_admin import Progress
from .types.datastore_admin import OperationType
from .types.index import Index
from .types.migration import MigrationProgressEvent
from .types.migration import MigrationStateEvent
from .types.migration import MigrationState
from .types.migration import MigrationStep

__all__ = (
"DatastoreAdminAsyncClient",
"CommonMetadata",
"CreateIndexRequest",
"DatastoreAdminClient",
"DatastoreFirestoreMigrationMetadata",
"DeleteIndexRequest",
"EntityFilter",
"ExportEntitiesMetadata",
Expand All @@ -51,6 +57,10 @@
"IndexOperationMetadata",
"ListIndexesRequest",
"ListIndexesResponse",
"MigrationProgressEvent",
"MigrationState",
"MigrationStateEvent",
"MigrationStep",
"OperationType",
"Progress",
)
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from collections import OrderedDict
import functools
import re
from typing import Dict, Sequence, Tuple, Type, Union
from typing import Dict, Optional, Sequence, Tuple, Type, Union
import pkg_resources

from google.api_core.client_options import ClientOptions
Expand All @@ -31,12 +31,12 @@
except AttributeError: # pragma: NO COVER
OptionalRetry = Union[retries.Retry, object] # type: ignore

from google.api_core import operation
from google.api_core import operation_async
from google.api_core import operation # type: ignore
from google.api_core import operation_async # type: ignore
from google.cloud.datastore_admin_v1.services.datastore_admin import pagers
from google.cloud.datastore_admin_v1.types import datastore_admin
from google.cloud.datastore_admin_v1.types import index
from google.protobuf import empty_pb2
from google.protobuf import empty_pb2 # type: ignore
from .transports.base import DatastoreAdminTransport, DEFAULT_CLIENT_INFO
from .transports.grpc_asyncio import DatastoreAdminGrpcAsyncIOTransport
from .client import DatastoreAdminClient
Expand Down Expand Up @@ -166,6 +166,42 @@ def from_service_account_file(cls, filename: str, *args, **kwargs):

from_service_account_json = from_service_account_file

@classmethod
def get_mtls_endpoint_and_cert_source(
cls, client_options: Optional[ClientOptions] = None
):
"""Return the API endpoint and client cert source for mutual TLS.

The client cert source is determined in the following order:
(1) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is not "true", the
client cert source is None.
(2) if `client_options.client_cert_source` is provided, use the provided one; if the
default client cert source exists, use the default one; otherwise the client cert
source is None.

The API endpoint is determined in the following order:
(1) if `client_options.api_endpoint` if provided, use the provided one.
(2) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is "always", use the
default mTLS endpoint; if the environment variabel is "never", use the default API
endpoint; otherwise if client cert source exists, use the default mTLS endpoint, otherwise
use the default API endpoint.

More details can be found at https://google.aip.dev/auth/4114.

Args:
client_options (google.api_core.client_options.ClientOptions): Custom options for the
client. Only the `api_endpoint` and `client_cert_source` properties may be used
in this method.

Returns:
Tuple[str, Callable[[], Tuple[bytes, bytes]]]: returns the API endpoint and the
client cert source to use.

Raises:
google.auth.exceptions.MutualTLSChannelError: If any errors happen.
"""
return DatastoreAdminClient.get_mtls_endpoint_and_cert_source(client_options) # type: ignore

@property
def transport(self) -> DatastoreAdminTransport:
"""Returns the transport used by the client instance.
Expand Down Expand Up @@ -316,7 +352,7 @@ async def export_entities(

"""
# Create or coerce a protobuf request object.
# Sanity check: If we got a request object, we should *not* have
# Quick check: If we got a request object, we should *not* have
# gotten any keyword arguments that map to the request.
has_flattened_params = any(
[project_id, labels, entity_filter, output_url_prefix]
Expand Down Expand Up @@ -456,7 +492,7 @@ async def import_entities(

"""
# Create or coerce a protobuf request object.
# Sanity check: If we got a request object, we should *not* have
# Quick check: If we got a request object, we should *not* have
# gotten any keyword arguments that map to the request.
has_flattened_params = any([project_id, labels, input_url, entity_filter])
if request is not None and has_flattened_params:
Expand Down
137 changes: 89 additions & 48 deletions google/cloud/datastore_admin_v1/services/datastore_admin/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,12 @@
except AttributeError: # pragma: NO COVER
OptionalRetry = Union[retries.Retry, object] # type: ignore

from google.api_core import operation
from google.api_core import operation_async
from google.api_core import operation # type: ignore
from google.api_core import operation_async # type: ignore
from google.cloud.datastore_admin_v1.services.datastore_admin import pagers
from google.cloud.datastore_admin_v1.types import datastore_admin
from google.cloud.datastore_admin_v1.types import index
from google.protobuf import empty_pb2
from google.protobuf import empty_pb2 # type: ignore
from .transports.base import DatastoreAdminTransport, DEFAULT_CLIENT_INFO
from .transports.grpc import DatastoreAdminGrpcTransport
from .transports.grpc_asyncio import DatastoreAdminGrpcAsyncIOTransport
Expand Down Expand Up @@ -277,6 +277,73 @@ def parse_common_location_path(path: str) -> Dict[str, str]:
m = re.match(r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)$", path)
return m.groupdict() if m else {}

@classmethod
def get_mtls_endpoint_and_cert_source(
cls, client_options: Optional[client_options_lib.ClientOptions] = None
):
"""Return the API endpoint and client cert source for mutual TLS.

The client cert source is determined in the following order:
(1) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is not "true", the
client cert source is None.
(2) if `client_options.client_cert_source` is provided, use the provided one; if the
default client cert source exists, use the default one; otherwise the client cert
source is None.

The API endpoint is determined in the following order:
(1) if `client_options.api_endpoint` if provided, use the provided one.
(2) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is "always", use the
default mTLS endpoint; if the environment variabel is "never", use the default API
endpoint; otherwise if client cert source exists, use the default mTLS endpoint, otherwise
use the default API endpoint.

More details can be found at https://google.aip.dev/auth/4114.

Args:
client_options (google.api_core.client_options.ClientOptions): Custom options for the
client. Only the `api_endpoint` and `client_cert_source` properties may be used
in this method.

Returns:
Tuple[str, Callable[[], Tuple[bytes, bytes]]]: returns the API endpoint and the
client cert source to use.

Raises:
google.auth.exceptions.MutualTLSChannelError: If any errors happen.
"""
if client_options is None:
client_options = client_options_lib.ClientOptions()
use_client_cert = os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false")
use_mtls_endpoint = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto")
if use_client_cert not in ("true", "false"):
raise ValueError(
"Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`"
)
if use_mtls_endpoint not in ("auto", "never", "always"):
raise MutualTLSChannelError(
"Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be `never`, `auto` or `always`"
)

# Figure out the client cert source to use.
client_cert_source = None
if use_client_cert == "true":
if client_options.client_cert_source:
client_cert_source = client_options.client_cert_source
elif mtls.has_default_client_cert_source():
client_cert_source = mtls.default_client_cert_source()

# Figure out which api endpoint to use.
if client_options.api_endpoint is not None:
api_endpoint = client_options.api_endpoint
elif use_mtls_endpoint == "always" or (
use_mtls_endpoint == "auto" and client_cert_source
):
api_endpoint = cls.DEFAULT_MTLS_ENDPOINT
else:
api_endpoint = cls.DEFAULT_ENDPOINT

return api_endpoint, client_cert_source

def __init__(
self,
*,
Expand Down Expand Up @@ -327,57 +394,22 @@ def __init__(
if client_options is None:
client_options = client_options_lib.ClientOptions()

# Create SSL credentials for mutual TLS if needed.
if os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false") not in (
"true",
"false",
):
raise ValueError(
"Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`"
)
use_client_cert = (
os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false") == "true"
api_endpoint, client_cert_source_func = self.get_mtls_endpoint_and_cert_source(
client_options
)

client_cert_source_func = None
is_mtls = False
if use_client_cert:
if client_options.client_cert_source:
is_mtls = True
client_cert_source_func = client_options.client_cert_source
else:
is_mtls = mtls.has_default_client_cert_source()
if is_mtls:
client_cert_source_func = mtls.default_client_cert_source()
else:
client_cert_source_func = None

# Figure out which api endpoint to use.
if client_options.api_endpoint is not None:
api_endpoint = client_options.api_endpoint
else:
use_mtls_env = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto")
if use_mtls_env == "never":
api_endpoint = self.DEFAULT_ENDPOINT
elif use_mtls_env == "always":
api_endpoint = self.DEFAULT_MTLS_ENDPOINT
elif use_mtls_env == "auto":
if is_mtls:
api_endpoint = self.DEFAULT_MTLS_ENDPOINT
else:
api_endpoint = self.DEFAULT_ENDPOINT
else:
raise MutualTLSChannelError(
"Unsupported GOOGLE_API_USE_MTLS_ENDPOINT value. Accepted "
"values: never, auto, always"
)
api_key_value = getattr(client_options, "api_key", None)
if api_key_value and credentials:
raise ValueError(
"client_options.api_key and credentials are mutually exclusive"
)

# Save or instantiate the transport.
# Ordinarily, we provide the transport, but allowing a custom transport
# instance provides an extensibility point for unusual situations.
if isinstance(transport, DatastoreAdminTransport):
# transport is a DatastoreAdminTransport instance.
if credentials or client_options.credentials_file:
if credentials or client_options.credentials_file or api_key_value:
raise ValueError(
"When providing a transport instance, "
"provide its credentials directly."
Expand All @@ -389,6 +421,15 @@ def __init__(
)
self._transport = transport
else:
import google.auth._default # type: ignore

if api_key_value and hasattr(
google.auth._default, "get_api_key_credentials"
):
credentials = google.auth._default.get_api_key_credentials(
api_key_value
)

Transport = type(self).get_transport_class(transport)
self._transport = Transport(
credentials=credentials,
Expand Down Expand Up @@ -491,7 +532,7 @@ def export_entities(

"""
# Create or coerce a protobuf request object.
# Sanity check: If we got a request object, we should *not* have
# Quick check: If we got a request object, we should *not* have
# gotten any keyword arguments that map to the request.
has_flattened_params = any(
[project_id, labels, entity_filter, output_url_prefix]
Expand Down Expand Up @@ -630,7 +671,7 @@ def import_entities(

"""
# Create or coerce a protobuf request object.
# Sanity check: If we got a request object, we should *not* have
# Quick check: If we got a request object, we should *not* have
# gotten any keyword arguments that map to the request.
has_flattened_params = any([project_id, labels, input_url, entity_filter])
if request is not None and has_flattened_params:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

from google.cloud.datastore_admin_v1.types import datastore_admin
from google.cloud.datastore_admin_v1.types import index
from google.longrunning import operations_pb2
from google.longrunning import operations_pb2 # type: ignore

try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down Expand Up @@ -107,7 +107,6 @@ def __init__(
credentials, _ = google.auth.load_credentials_from_file(
credentials_file, **scopes_kwargs, quota_project_id=quota_project_id
)

elif credentials is None:
credentials, _ = google.auth.default(
**scopes_kwargs, quota_project_id=quota_project_id
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,11 @@
from google.auth import credentials as ga_credentials # type: ignore
from google.auth.transport.grpc import SslCredentials # type: ignore

import grpc
import grpc # type: ignore

from google.cloud.datastore_admin_v1.types import datastore_admin
from google.cloud.datastore_admin_v1.types import index
from google.longrunning import operations_pb2
from google.longrunning import operations_pb2 # type: ignore
from .base import DatastoreAdminTransport, DEFAULT_CLIENT_INFO


Expand Down Expand Up @@ -219,8 +219,11 @@ def __init__(
if not self._grpc_channel:
self._grpc_channel = type(self).create_channel(
self._host,
# use the credentials which are saved
credentials=self._credentials,
credentials_file=credentials_file,
# Set ``credentials_file`` to ``None`` here as
# the credentials that we saved earlier should be used.
credentials_file=None,
scopes=self._scopes,
ssl_credentials=self._ssl_channel_credentials,
quota_project_id=quota_project_id,
Expand Down Expand Up @@ -293,7 +296,7 @@ def operations_client(self) -> operations_v1.OperationsClient:
This property caches on the instance; repeated calls return the same
client.
"""
# Sanity check: Only create a new client if we do not already have one.
# Quick check: Only create a new client if we do not already have one.
if self._operations_client is None:
self._operations_client = operations_v1.OperationsClient(self.grpc_channel)

Expand Down
Loading