def setUpModule():
from google.cloud.exceptions import GrpcRendezvous
Config.IN_EMULATOR = os.getenv(BIGTABLE_EMULATOR) is not None
if Config.IN_EMULATOR:
credentials = EmulatorCreds()
Config.CLIENT = Client(admin=True, credentials=credentials)
else:
Config.CLIENT = Client(admin=True)
Config.INSTANCE = Config.CLIENT.instance(INSTANCE_ID, LOCATION_ID)
if not Config.IN_EMULATOR:
retry = RetryErrors(GrpcRendezvous,
error_predicate=_retry_on_unavailable)
instances_response = retry(Config.CLIENT.list_instances)()
if len(instances_response.failed_locations) != 0:
raise ValueError('List instances failed in module set up.')
EXISTING_INSTANCES[:] = instances_response.instances
# After listing, create the test instance.
created_op = Config.INSTANCE.create()
created_op.result(timeout=10)
def _get_instance():
"""Gets instance for the default project.
Creates a client with the inferred credentials and project ID from
the local environment. Then uses
:meth:`.bigtable.client.Client.list_instances` to
get the unique instance owned by the project.
If the request fails for any reason, or if there isn't exactly one instance
owned by the project, then this function will fail.
:rtype: :class:`~google.cloud.bigtable.instance.Instance`
:returns: The unique instance owned by the project inferred from
the environment.
:raises ValueError: if there is a failed location or any number of
instances other than one.
"""
client_kwargs = {'admin': True}
client = Client(**client_kwargs)
instances, failed_locations = client.list_instances()
if failed_locations:
raise ValueError('Determining instance via ListInstances encountered '
'failed locations.')
num_instances = len(instances)
if num_instances == 0:
raise ValueError('This client doesn\'t have access to any instances.')
if num_instances > 1:
raise ValueError('This client has access to more than one instance. '
'Please directly pass the instance you\'d '
'like to use.')
return instances[0]
def build_bigtable_client(project_id=None, credentials=None,
service_key_file=None, read_only=False):
"""
Build a Google Cloud Bigtable API client.
Args:
project_id (:obj:`str`, optional): The ID of the project owning the
instances and tables. Defaults to None.
credentials (:obj:`object`, optional): Credentials object to pass into
client. If None, attempt to determine from environment. Defaults to
None.
service_key_file (:obj:`str`, optional): Path to a JSON keyfile
containing service account credentials. This is recommended for
production environments. Defaults to None.
read_only (:obj:`bool`, optional): Indicates if the client should be
for reading only. Defaults to False.
Returns:
Client: Google Cloud Bigtable API client to work with.
"""
admin = not read_only # Cannot perform admin tasks if read-only.
kwargs = {'admin': admin, 'read_only': read_only}
if service_key_file and os.path.isfile(service_key_file):
client = Client.from_service_account_json(service_key_file, **kwargs)
else:
client = Client(project=project_id, credentials=credentials, **kwargs)
return client
def test_test_iam_permissions(self):
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import (
BigtableTableAdminClient,
)
from google.iam.v1 import iam_policy_pb2
credentials = _make_credentials()
client = Client(project=self.PROJECT_ID, credentials=credentials, admin=True)
instance = client.instance(instance_id=self.INSTANCE_ID)
backup = self._make_one(self.BACKUP_ID, instance, cluster_id=self.CLUSTER_ID)
permissions = ["bigtable.backups.create", "bigtable.backups.list"]
response = iam_policy_pb2.TestIamPermissionsResponse(permissions=permissions)
table_api = mock.create_autospec(BigtableTableAdminClient)
table_api.test_iam_permissions.return_value = response
client._table_admin_client = table_api
result = backup.test_iam_permissions(permissions)
self.assertEqual(result, permissions)
table_api.test_iam_permissions.assert_called_once_with(
request={"resource": backup.name, "permissions": permissions}
)
def setUpModule():
from google.cloud.exceptions import GrpcRendezvous
Config.IN_EMULATOR = os.getenv(BIGTABLE_EMULATOR) is not None
if Config.IN_EMULATOR:
credentials = EmulatorCreds()
Config.CLIENT = Client(admin=True, credentials=credentials)
else:
Config.CLIENT = Client(admin=True)
Config.INSTANCE = Config.CLIENT.instance(INSTANCE_ID, labels=LABELS)
Config.CLUSTER = Config.INSTANCE.cluster(CLUSTER_ID,
location_id=LOCATION_ID,
serve_nodes=SERVE_NODES)
if not Config.IN_EMULATOR:
retry = RetryErrors(GrpcRendezvous,
error_predicate=_retry_on_unavailable)
instances, failed_locations = retry(Config.CLIENT.list_instances)()
if len(failed_locations) != 0:
raise ValueError("List instances failed in module set up.")
EXISTING_INSTANCES[:] = instances
# After listing, create the test instance.
created_op = Config.INSTANCE.create(clusters=[Config.CLUSTER])
created_op.result(timeout=10)
def __init__(self, PROJECT_ID, INSTANCE_ID, TABLE_ID):
self.project_id = PROJECT_ID
self.instance_id = INSTANCE_ID
self.table_id = TABLE_ID
client = Client(project=PROJECT_ID, admin=True)
instance = client.instance(INSTANCE_ID)
table = instance.table(TABLE_ID)
def test_property_source_table_valid(self):
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client
api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock())
credentials = _make_credentials()
client = Client(project=self.PROJECT_ID, credentials=credentials, admin=True)
client._table_admin_client = api
instance = _Instance(self.INSTANCE_NAME, client)
backup = self._make_one(self.BACKUP_ID, instance, table_id=self.TABLE_ID)
self.assertEqual(backup.source_table, self.TABLE_NAME)
def test_set_iam_policy(self):
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import (
BigtableTableAdminClient, )
from google.iam.v1 import policy_pb2
from google.cloud.bigtable.policy import Policy
from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE
credentials = _make_credentials()
client = Client(project=self.PROJECT_ID,
credentials=credentials,
admin=True)
instance = client.instance(instance_id=self.INSTANCE_ID)
backup = self._make_one(self.BACKUP_ID,
instance,
cluster_id=self.CLUSTER_ID)
version = 1
etag = b"etag_v1"
members = [
"serviceAccount:[email protected]", "user:[email protected]"
]
bindings = [{"role": BIGTABLE_ADMIN_ROLE, "members": sorted(members)}]
iam_policy_pb = policy_pb2.Policy(version=version,
etag=etag,
bindings=bindings)
table_api = mock.create_autospec(BigtableTableAdminClient)
client._table_admin_client = table_api
table_api.set_iam_policy.return_value = iam_policy_pb
iam_policy = Policy(etag=etag, version=version)
iam_policy[BIGTABLE_ADMIN_ROLE] = [
Policy.user("*****@*****.**"),
Policy.service_account("*****@*****.**"),
]
result = backup.set_iam_policy(iam_policy)
table_api.set_iam_policy.assert_called_once_with(
request={
"resource": backup.name,
"policy": iam_policy_pb
})
self.assertEqual(result.version, version)
self.assertEqual(result.etag, etag)
admins = result.bigtable_admins
self.assertEqual(len(admins), len(members))
for found, expected in zip(sorted(admins), sorted(members)):
self.assertEqual(found, expected)
def test_backup_source_table_valid():
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient, )
api = mock.create_autospec(BigtableInstanceAdminClient)
credentials = _make_credentials()
client = Client(project=PROJECT_ID, credentials=credentials, admin=True)
client._table_admin_client = api
instance = _Instance(INSTANCE_NAME, client)
backup = _make_backup(BACKUP_ID, instance, table_id=TABLE_ID)
assert backup.source_table == TABLE_NAME
def test_property_parent_w_cluster(self):
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client
api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock())
credentials = _make_credentials()
client = Client(project=self.PROJECT_ID, credentials=credentials, admin=True)
client._table_admin_client = api
instance = _Instance(self.INSTANCE_NAME, client)
backup = self._make_one(self.BACKUP_ID, instance, cluster_id=self.CLUSTER_ID)
self.assertEqual(backup._cluster, self.CLUSTER_ID)
self.assertEqual(backup.parent, self.CLUSTER_NAME)
def test_property_source_table_valid(self):
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient,
)
api = mock.create_autospec(BigtableInstanceAdminClient)
credentials = _make_credentials()
client = Client(project=self.PROJECT_ID, credentials=credentials, admin=True)
client._table_admin_client = api
instance = _Instance(self.INSTANCE_NAME, client)
backup = self._make_one(self.BACKUP_ID, instance, table_id=self.TABLE_ID)
self.assertEqual(backup.source_table, self.TABLE_NAME)
def open(self):
self._table = Client(project=self._config.project,
credentials=self._config.credentials,
admin=True).instance(self._config.instance).table(
self._config.table)
# maybe add batcher later as an optimization?
# self._buffer = self._table.mutations_batcher()
# ensure table and column are created
if not self._table.exists():
self._table.create()
columns = self._table.list_column_families().keys()
if FAMILY not in columns:
self._table.column_family(FAMILY).create()
return self
def test_backup_parent_w_cluster():
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient, )
api = mock.create_autospec(BigtableInstanceAdminClient)
credentials = _make_credentials()
client = Client(project=PROJECT_ID, credentials=credentials, admin=True)
client._table_admin_client = api
instance = _Instance(INSTANCE_NAME, client)
backup = _make_backup(BACKUP_ID, instance, cluster_id=CLUSTER_ID)
assert backup._cluster == CLUSTER_ID
assert backup.parent == CLUSTER_NAME
def test_backup_create_success():
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.bigtable_admin_v2.types import table
from google.cloud.bigtable import Client
op_future = object()
credentials = _make_credentials()
client = Client(project=PROJECT_ID, credentials=credentials, admin=True)
api = client._table_admin_client = _make_table_admin_client()
api.create_backup.return_value = op_future
timestamp = _make_timestamp()
backup = _make_backup(
BACKUP_ID,
_Instance(INSTANCE_NAME, client=client),
table_id=TABLE_ID,
expire_time=timestamp,
)
backup_pb = table.Backup(
source_table=TABLE_NAME,
expire_time=_datetime_to_pb_timestamp(timestamp),
)
future = backup.create(CLUSTER_ID)
assert backup._cluster == CLUSTER_ID
assert future is op_future
api.create_backup.assert_called_once_with(request={
"parent": CLUSTER_NAME,
"backup_id": BACKUP_ID,
"backup": backup_pb
})
def test_create_success(self):
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.bigtable_admin_v2.types import table
from google.cloud.bigtable import Client
op_future = object()
credentials = _make_credentials()
client = Client(project=self.PROJECT_ID, credentials=credentials, admin=True)
api = client._table_admin_client = self._make_table_admin_client()
api.create_backup.return_value = op_future
timestamp = self._make_timestamp()
backup = self._make_one(
self.BACKUP_ID,
_Instance(self.INSTANCE_NAME, client=client),
table_id=self.TABLE_ID,
expire_time=timestamp,
)
backup_pb = table.Backup(
source_table=self.TABLE_NAME,
expire_time=_datetime_to_pb_timestamp(timestamp),
)
future = backup.create(self.CLUSTER_ID)
self.assertEqual(backup._cluster, self.CLUSTER_ID)
self.assertIs(future, op_future)
api.create_backup.assert_called_once_with(
request={
"parent": self.CLUSTER_NAME,
"backup_id": self.BACKUP_ID,
"backup": backup_pb,
}
)
def setUpModule():
from google.cloud.exceptions import GrpcRendezvous
from google.cloud.bigtable.enums import Instance
# See: https://github.com/googleapis/google-cloud-python/issues/5928
interfaces = table_admin_config.config["interfaces"]
iface_config = interfaces["google.bigtable.admin.v2.BigtableTableAdmin"]
methods = iface_config["methods"]
create_table = methods["CreateTable"]
create_table["timeout_millis"] = 90000
Config.IN_EMULATOR = os.getenv(BIGTABLE_EMULATOR) is not None
if Config.IN_EMULATOR:
credentials = EmulatorCreds()
Config.CLIENT = Client(admin=True, credentials=credentials)
else:
Config.CLIENT = Client(admin=True)
Config.INSTANCE = Config.CLIENT.instance(INSTANCE_ID, labels=LABELS)
Config.CLUSTER = Config.INSTANCE.cluster(CLUSTER_ID,
location_id=LOCATION_ID,
serve_nodes=SERVE_NODES)
Config.INSTANCE_DATA = Config.CLIENT.instance(
INSTANCE_ID_DATA,
instance_type=Instance.Type.DEVELOPMENT,
labels=LABELS)
Config.CLUSTER_DATA = Config.INSTANCE_DATA.cluster(CLUSTER_ID_DATA,
location_id=LOCATION_ID)
if not Config.IN_EMULATOR:
retry = RetryErrors(GrpcRendezvous,
error_predicate=_retry_on_unavailable)
instances, failed_locations = retry(Config.CLIENT.list_instances)()
if len(failed_locations) != 0:
raise ValueError("List instances failed in module set up.")
EXISTING_INSTANCES[:] = instances
# After listing, create the test instances.
admin_op = Config.INSTANCE.create(clusters=[Config.CLUSTER])
admin_op.result(timeout=10)
data_op = Config.INSTANCE_DATA.create(clusters=[Config.CLUSTER_DATA])
data_op.result(timeout=10)
class BigTableMap(Map):
def open(self):
self._table = Client(project=self._config.project,
credentials=self._config.credentials,
admin=True).instance(self._config.instance).table(
self._config.table)
# maybe add batcher later as an optimization?
# self._buffer = self._table.mutations_batcher()
# ensure table and column are created
if not self._table.exists():
self._table.create()
columns = self._table.list_column_families().keys()
if FAMILY not in columns:
self._table.column_family(FAMILY).create()
return self
def close(self):
"""does bigtable client require cleanup?"""
# maybe add batcher later as an optimization?
# self._buffer.flush()
return self
def _put(self, key: Key, value: bytes) -> Key:
row = self._table.row(row_key=key.uuid.bytes_le)
row.set_cell(column_family_id=FAMILY,
column=COLUMN,
value=value,
timestamp=millis_dt(key.time))
self._table.mutate_rows([row])
return key
def _get(self, uuid: UUID, time: int) -> Optional[bytes]:
return self._table.read_row(
row_key=uuid.bytes_le,
filter_=RowFilterChain(filters=[
TimestampRangeFilter(range_=TimestampRange(end=millis_dt(time +
1))),
CellsColumnLimitFilter(num_cells=1)
])).cell_value(column_family_id=FAMILY, column=COLUMN)
def test_backup_get_iam_policy():
from google.cloud.bigtable.client import Client
from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import (
BigtableTableAdminClient, )
from google.iam.v1 import policy_pb2
from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE
credentials = _make_credentials()
client = Client(project=PROJECT_ID, credentials=credentials, admin=True)
instance = client.instance(instance_id=INSTANCE_ID)
backup = _make_backup(BACKUP_ID, instance, cluster_id=CLUSTER_ID)
version = 1
etag = b"etag_v1"
members = ["serviceAccount:[email protected]", "user:[email protected]"]
bindings = [{"role": BIGTABLE_ADMIN_ROLE, "members": members}]
iam_policy = policy_pb2.Policy(version=version,
etag=etag,
bindings=bindings)
table_api = mock.create_autospec(BigtableTableAdminClient)
client._table_admin_client = table_api
table_api.get_iam_policy.return_value = iam_policy
result = backup.get_iam_policy()
table_api.get_iam_policy.assert_called_once_with(
request={"resource": backup.name})
assert result.version == version
assert result.etag == etag
admins = result.bigtable_admins
assert len(admins) == len(members)
for found, expected in zip(sorted(admins), sorted(members)):
assert found == expected
def setUpModule():
from grpc._channel import _Rendezvous
_helpers.PROJECT = TESTS_PROJECT
Config.CLIENT = Client(admin=True)
Config.INSTANCE = Config.CLIENT.instance(INSTANCE_ID, LOCATION_ID)
retry = RetryErrors(_Rendezvous, error_predicate=_retry_on_unavailable)
instances, failed_locations = retry(Config.CLIENT.list_instances)()
if len(failed_locations) != 0:
raise ValueError('List instances failed in module set up.')
EXISTING_INSTANCES[:] = instances
# After listing, create the test instance.
created_op = Config.INSTANCE.create()
if not _wait_until_complete(created_op):
raise RuntimeError('Instance creation exceed 5 seconds.')
def get_client(project_id, read_only=False):
"""
Get the client associated with the Bigtable emulator.
Args:
project_id (str): Project ID to work with. Any arbitrary string
will work.
read_only (:obj:`bool`, optional): Indicates whether client should
be read-only or not. Defaults to False.
Returns:
Client: Google Cloud Bigtable API client initialized with emulator
credentials. None if emulator is not running.
"""
admin = not read_only
client = Client(
project=project_id, credentials=EmulatorCredentials(),
admin=admin, read_only=read_only
)
return client if client.emulator_host is not None else None
def test_access_with_non_admin_client(self):
client = Client(admin=False)
instance = client.instance(INSTANCE_ID)
table = instance.table(self._table.table_id)
self.assertIsNone(table.read_row("nonesuch"))
def test_retry(self):
import subprocess, os, stat, platform
from google.cloud.bigtable.client import Client
from google.cloud.bigtable.instance import Instance
from google.cloud.bigtable.table import Table
# import for urlopen based on version
try:
# python 3
from urllib.request import urlopen
except ImportError:
# python 2
from urllib2 import urlopen
TEST_SCRIPT = 'tests/retry_test_script.txt'
SERVER_NAME = 'retry_server'
SERVER_ZIP = SERVER_NAME + ".tar.gz"
def process_scan(table, range, ids):
range_chunks = range.split(",")
range_open = range_chunks[0].lstrip("[")
range_close = range_chunks[1].rstrip(")")
rows = table.read_rows(range_open, range_close)
rows.consume_all()
# Download server
MOCK_SERVER_URLS = {
'Linux':
'https://storage.googleapis.com/cloud-bigtable-test/retries/retry_server_linux.tar.gz',
'Darwin':
'https://storage.googleapis.com/cloud-bigtable-test/retries/retry_server_mac.tar.gz',
}
test_platform = platform.system()
if test_platform not in MOCK_SERVER_URLS:
self.skip('Retry server not available for platform {0}.'.format(
test_platform))
mock_server_download = urlopen(MOCK_SERVER_URLS[test_platform]).read()
mock_server_file = open(SERVER_ZIP, 'wb')
mock_server_file.write(mock_server_download)
# Unzip server
subprocess.call(['tar', 'zxvf', SERVER_ZIP, '-C', '.'])
# Connect to server
server = subprocess.Popen(
['./' + SERVER_NAME, '--script=' + TEST_SCRIPT],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
(endpoint, port) = server.stdout.readline().rstrip("\n").split(":")
os.environ["BIGTABLE_EMULATOR_HOST"] = endpoint + ":" + port
client = Client(project="client", admin=True)
instance = Instance("instance", client)
table = instance.table("table")
# Run test, line by line
with open(TEST_SCRIPT, 'r') as script:
for line in script.readlines():
if line.startswith("CLIENT:"):
chunks = line.split(" ")
op = chunks[1]
process_scan(table, chunks[2], chunks[3])
# Check that the test passed
server.kill()
server_stdout_lines = []
while True:
line = server.stdout.readline()
if line != '':
server_stdout_lines.append(line)
else:
break
self.assertEqual(server_stdout_lines[-1], "PASS\n")
# Clean up
os.remove(SERVER_ZIP)
os.remove(SERVER_NAME)
def data_client():
return Client(admin=False)
def admin_client():
return Client(admin=True)
def _make_client(*args, **kwargs):
from google.cloud.bigtable.client import Client
return Client(*args, **kwargs)
def test_bigtable_client(self):
# Sanity test to ensure that we are using the emulator for testing.
client = Client(PROJECT_ID, credentials=EmulatorCredentials())
self.assertIn('localhost', client.emulator_host)
def test_access_with_non_admin_client(self):
client = Client(admin=False)
instance = client.instance(INSTANCE_ID)
table = instance.table(self._table.table_id)
self.assertIsNone(table.read_row("nonesuch"))
