387 lines
14 KiB
Python
387 lines
14 KiB
Python
#!/usr/bin/env python
|
|
"""BQ CLI helper functions for IDs."""
|
|
|
|
import collections
|
|
import sys
|
|
from typing import Any, Optional, Tuple, Type, Union
|
|
from absl import app
|
|
from utils import bq_error
|
|
from pyglib import stringutil
|
|
|
|
collections_abc = collections
|
|
if sys.version_info > (3, 8):
|
|
collections_abc = collections.abc
|
|
|
|
|
|
class ApiClientHelper:
|
|
"""Static helper methods and classes not provided by the discovery client."""
|
|
|
|
def __init__(self, *unused_args, **unused_kwds):
|
|
raise NotImplementedError('Cannot instantiate static class ApiClientHelper')
|
|
|
|
class Reference(collections_abc.Mapping):
|
|
"""Base class for Reference objects returned by apiclient."""
|
|
|
|
_required_fields = frozenset()
|
|
_optional_fields = frozenset()
|
|
_format_str = ''
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=unidiomatic-typecheck Check if this isn't a subclass.
|
|
if type(self) == ApiClientHelper.Reference:
|
|
self.typename: str = 'unimplemented'
|
|
raise NotImplementedError(
|
|
'Cannot instantiate abstract class ApiClientHelper.Reference'
|
|
)
|
|
for name in self._required_fields:
|
|
if not kwds.get(name, ''):
|
|
raise ValueError(
|
|
'Missing required argument %s to %s'
|
|
% (name, self.__class__.__name__)
|
|
)
|
|
setattr(self, name, kwds[name])
|
|
for name in self._optional_fields:
|
|
if kwds.get(name, ''):
|
|
setattr(self, name, kwds[name])
|
|
|
|
@classmethod
|
|
def Create(cls, **kwds: Any) -> 'ApiClientHelper.Reference':
|
|
"""Factory method for this class."""
|
|
args = dict(
|
|
(k, v)
|
|
for k, v in kwds.items()
|
|
if k in cls._required_fields.union(cls._optional_fields)
|
|
)
|
|
return cls(**args)
|
|
|
|
def __iter__(self):
|
|
return iter(self._required_fields.union(self._optional_fields))
|
|
|
|
def __getitem__(self, key):
|
|
if key in self._optional_fields:
|
|
if key in self.__dict__:
|
|
return self.__dict__[key]
|
|
else:
|
|
return None
|
|
if key in self._required_fields:
|
|
return self.__dict__[key]
|
|
raise KeyError(key)
|
|
|
|
def __hash__(self):
|
|
return hash(str(self))
|
|
|
|
def __len__(self):
|
|
return len(self._required_fields.union(self._optional_fields))
|
|
|
|
def __str__(self):
|
|
return stringutil.ensure_str(self._format_str % dict(self))
|
|
|
|
def __repr__(self):
|
|
return "%s '%s'" % (self.typename, self)
|
|
|
|
def __eq__(self, other):
|
|
d = dict(other)
|
|
return all(
|
|
getattr(self, name, None) == d.get(name, None)
|
|
for name in self._required_fields.union(self._optional_fields)
|
|
)
|
|
|
|
class JobReference(Reference):
|
|
"""A JobReference."""
|
|
|
|
_required_fields = frozenset(('projectId', 'jobId'))
|
|
_optional_fields = frozenset(('location',))
|
|
_format_str = '%(projectId)s:%(jobId)s'
|
|
typename = 'job'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.jobId: str = kwds['jobId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def GetProjectReference(self) -> 'ApiClientHelper.ProjectReference':
|
|
return ApiClientHelper.ProjectReference.Create(projectId=self.projectId)
|
|
|
|
class ProjectReference(Reference):
|
|
"""A ProjectReference."""
|
|
|
|
_required_fields = frozenset(('projectId',))
|
|
_format_str = '%(projectId)s'
|
|
typename = 'project'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def GetDatasetReference(
|
|
self, dataset_id: str
|
|
) -> 'ApiClientHelper.DatasetReference':
|
|
return ApiClientHelper.DatasetReference.Create(
|
|
projectId=self.projectId, datasetId=dataset_id
|
|
)
|
|
|
|
def GetTableReference(
|
|
self, dataset_id: str, table_id: str
|
|
) -> 'ApiClientHelper.TableReference':
|
|
return ApiClientHelper.TableReference.Create(
|
|
projectId=self.projectId, datasetId=dataset_id, tableId=table_id
|
|
)
|
|
|
|
class DatasetReference(Reference):
|
|
"""A DatasetReference."""
|
|
|
|
_required_fields = frozenset(('projectId', 'datasetId'))
|
|
_format_str = '%(projectId)s:%(datasetId)s'
|
|
typename = 'dataset'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.datasetId: str = kwds['datasetId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def GetProjectReference(self) -> 'ApiClientHelper.ProjectReference':
|
|
return ApiClientHelper.ProjectReference.Create(projectId=self.projectId)
|
|
|
|
def GetTableReference(
|
|
self, table_id: str
|
|
) -> 'ApiClientHelper.TableReference':
|
|
return ApiClientHelper.TableReference.Create(
|
|
projectId=self.projectId, datasetId=self.datasetId, tableId=table_id
|
|
)
|
|
|
|
|
|
class TableReference(Reference):
|
|
"""A TableReference."""
|
|
|
|
_required_fields = frozenset(('projectId', 'datasetId', 'tableId'))
|
|
_format_str = '%(projectId)s:%(datasetId)s.%(tableId)s'
|
|
typename = 'table'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.datasetId: str = kwds['datasetId']
|
|
self.tableId: str = kwds['tableId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def GetDatasetReference(self) -> 'ApiClientHelper.DatasetReference':
|
|
return ApiClientHelper.DatasetReference.Create(
|
|
projectId=self.projectId, datasetId=self.datasetId
|
|
)
|
|
|
|
def GetProjectReference(self) -> 'ApiClientHelper.ProjectReference':
|
|
return ApiClientHelper.ProjectReference.Create(projectId=self.projectId)
|
|
|
|
class ModelReference(Reference):
|
|
_required_fields = frozenset(('projectId', 'datasetId', 'modelId'))
|
|
_format_str = '%(projectId)s:%(datasetId)s.%(modelId)s'
|
|
typename = 'model'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.datasetId: str = kwds['datasetId']
|
|
self.modelId: str = kwds['modelId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
class RoutineReference(Reference):
|
|
"""A RoutineReference."""
|
|
|
|
_required_fields = frozenset(('projectId', 'datasetId', 'routineId'))
|
|
_format_str = '%(projectId)s:%(datasetId)s.%(routineId)s'
|
|
_path_str = (
|
|
'projects/%(projectId)s/datasets/%(datasetId)s/routines/%(routineId)s'
|
|
)
|
|
|
|
typename = 'routine'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.datasetId: str = kwds['datasetId']
|
|
self.routineId: str = kwds['routineId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def path(self) -> str:
|
|
return self._path_str % dict(self)
|
|
|
|
class RowAccessPolicyReference(Reference):
|
|
_required_fields = frozenset(
|
|
('projectId', 'datasetId', 'tableId', 'policyId')
|
|
)
|
|
_format_str = '%(projectId)s:%(datasetId)s.%(tableId)s.%(policyId)s'
|
|
typename = 'row access policy'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.datasetId: str = kwds['datasetId']
|
|
self.tableId: str = kwds['tableId']
|
|
self.policyId: str = kwds['policyId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
class TransferConfigReference(Reference):
|
|
_required_fields = frozenset(('transferConfigName',))
|
|
_format_str = '%(transferConfigName)s'
|
|
typename = 'transfer config'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.transferConfigName: str = kwds['transferConfigName']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
class TransferRunReference(Reference):
|
|
_required_fields = frozenset(('transferRunName',))
|
|
_format_str = '%(transferRunName)s'
|
|
typename = 'transfer run'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.transferRunName: str = kwds['transferRunName']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
class NextPageTokenReference(Reference):
|
|
_required_fields = frozenset(('pageTokenId',))
|
|
_format_str = '%(pageTokenId)s'
|
|
typename = 'page token'
|
|
|
|
class TransferLogReference(TransferRunReference):
|
|
pass
|
|
|
|
class EncryptionServiceAccount(Reference):
|
|
_required_fields = frozenset(('serviceAccount',))
|
|
_format_str = '%(serviceAccount)s'
|
|
# typename is set to none because the EncryptionServiceAccount does not
|
|
# store a 'reference', so when the object info is printed, it will omit
|
|
# an unnecessary line that would have tried to print a reference in other
|
|
# cases, i.e. datasets, tables, etc.
|
|
typename = None
|
|
|
|
class ReservationReference(Reference):
|
|
_required_fields = frozenset(('projectId', 'location', 'reservationId'))
|
|
_format_str = '%(projectId)s:%(location)s.%(reservationId)s'
|
|
_path_str = 'projects/%(projectId)s/locations/%(location)s/reservations/%(reservationId)s'
|
|
typename = 'reservation'
|
|
|
|
def path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._path_str % dict(self)
|
|
|
|
class CapacityCommitmentReference(Reference):
|
|
"""Helper class to provide a reference to capacity commitment."""
|
|
|
|
_required_fields = frozenset(
|
|
('projectId', 'location', 'capacityCommitmentId')
|
|
)
|
|
_format_str = '%(projectId)s:%(location)s.%(capacityCommitmentId)s'
|
|
_path_str = 'projects/%(projectId)s/locations/%(location)s/capacityCommitments/%(capacityCommitmentId)s'
|
|
typename = 'capacity commitment'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.location: str = kwds['location']
|
|
self.capacityCommitmentId: str = kwds['capacityCommitmentId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._path_str % dict(self)
|
|
|
|
class ReservationAssignmentReference(Reference):
|
|
"""Helper class to provide a reference to reservation assignment."""
|
|
|
|
_required_fields = frozenset(
|
|
('projectId', 'location', 'reservationId', 'reservationAssignmentId')
|
|
)
|
|
_format_str = '%(projectId)s:%(location)s.%(reservationId)s.%(reservationAssignmentId)s'
|
|
_path_str = 'projects/%(projectId)s/locations/%(location)s/reservations/%(reservationId)s/assignments/%(reservationAssignmentId)s'
|
|
_reservation_format_str = '%(projectId)s:%(location)s.%(reservationId)s'
|
|
typename = 'reservation assignment'
|
|
|
|
def path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._path_str % dict(self)
|
|
|
|
def reservation_path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._reservation_format_str % dict(self)
|
|
|
|
class BiReservationReference(Reference):
|
|
"""Helper class to provide a reference to bi reservation."""
|
|
|
|
_required_fields = frozenset(('projectId', 'location'))
|
|
_format_str = '%(projectId)s:%(location)s'
|
|
_path_str = 'projects/%(projectId)s/locations/%(location)s/biReservation'
|
|
_create_path_str = 'projects/%(projectId)s/locations/%(location)s'
|
|
typename = 'bi reservation'
|
|
|
|
def path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._path_str % dict(self)
|
|
|
|
def create_path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._create_path_str % dict(self)
|
|
|
|
class ReservationGroupReference(Reference):
|
|
"""Helper class to provide a reference to reservation group."""
|
|
|
|
_required_fields = frozenset(
|
|
('projectId', 'location', 'reservationGroupId')
|
|
)
|
|
_format_str = '%(projectId)s:%(location)s.%(reservationGroupId)s'
|
|
_path_str = 'projects/%(projectId)s/locations/%(location)s/reservationGroups/%(reservationGroupId)s'
|
|
typename = 'reservation group'
|
|
|
|
def __init__(self, **kwds):
|
|
# pylint: disable=invalid-name Aligns with API
|
|
self.projectId: str = kwds['projectId']
|
|
self.location: str = kwds['location']
|
|
self.reservationGroupId: str = kwds['reservationGroupId']
|
|
# pylint: enable=invalid-name
|
|
super().__init__(**kwds)
|
|
|
|
def path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._path_str % dict(self)
|
|
|
|
class ConnectionReference(Reference):
|
|
_required_fields = frozenset(('projectId', 'location', 'connectionId'))
|
|
_format_str = '%(projectId)s.%(location)s.%(connectionId)s'
|
|
_path_str = 'projects/%(projectId)s/locations/%(location)s/connections/%(connectionId)s'
|
|
typename = 'connection'
|
|
|
|
def path(self) -> str: # pylint: disable=invalid-name Legacy
|
|
return self._path_str % dict(self)
|
|
|
|
|
|
def typecheck( # pylint: disable=invalid-name
|
|
obj: ApiClientHelper.Reference,
|
|
types: Union[
|
|
Type[Optional[ApiClientHelper.Reference]],
|
|
Tuple[Type[Optional[ApiClientHelper.Reference]], ...],
|
|
],
|
|
message: Optional[str] = None,
|
|
method: Optional[str] = None,
|
|
# In code on the surface, taking user input, we throw a usage error.
|
|
is_usage_error: bool = False,
|
|
) -> None:
|
|
"""Ensure the obj is the correct type, or throw a BigqueryTypeError."""
|
|
if not isinstance(obj, types):
|
|
if not message:
|
|
if method:
|
|
message = 'Invalid reference for %s: %r' % (method, obj)
|
|
else:
|
|
message = 'Type of %r is not one of %s' % (obj, types)
|
|
if is_usage_error:
|
|
raise app.UsageError(message)
|
|
else:
|
|
raise bq_error.BigqueryTypeError(message)
|