# Copyright 2016 Google LLC All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import unittest

import mock

from google.cloud.spanner_v1.param_types import INT64

DML_WO_PARAM = """

DELETE FROM citizens

"""

DML_W_PARAM = """

INSERT INTO citizens(first_name, last_name, age)

VALUES ("Phred", "Phlyntstone", @age)

"""

PARAMS = {"age": 30}

PARAM_TYPES = {"age": INT64}

MODE = 2 # PROFILE

def _make_credentials(): # pragma: NO COVER

import google.auth.credentials

class _CredentialsWithScopes(

google.auth.credentials.Credentials, google.auth.credentials.Scoped

):

pass

return mock.Mock(spec=_CredentialsWithScopes)

class _BaseTest(unittest.TestCase):

PROJECT_ID = "project-id"

PARENT = "projects/" + PROJECT_ID

INSTANCE_ID = "instance-id"

INSTANCE_NAME = PARENT + "/instances/" + INSTANCE_ID

DATABASE_ID = "database_id"

DATABASE_NAME = INSTANCE_NAME + "/databases/" + DATABASE_ID

SESSION_ID = "session_id"

SESSION_NAME = DATABASE_NAME + "/sessions/" + SESSION_ID

TRANSACTION_ID = b"transaction_id"

RETRY_TRANSACTION_ID = b"transaction_id_retry"

BACKUP_ID = "backup_id"

BACKUP_NAME = INSTANCE_NAME + "/backups/" + BACKUP_ID

def _make_one(self, *args, **kwargs):

return self._get_target_class()(*args, **kwargs)

@staticmethod

def _make_timestamp():

import datetime

from google.cloud._helpers import UTC

return datetime.datetime.utcnow().replace(tzinfo=UTC)

@staticmethod

def _make_duration(seconds=1, microseconds=0):

import datetime

return datetime.timedelta(seconds=seconds, microseconds=microseconds)

class TestDatabase(_BaseTest):

def _get_target_class(self):

from google.cloud.spanner_v1.database import Database

return Database

@staticmethod

def _make_database_admin_api():

from google.cloud.spanner_v1.client import DatabaseAdminClient

return mock.create_autospec(DatabaseAdminClient, instance=True)

@staticmethod

def _make_spanner_api():

from google.cloud.spanner_v1 import SpannerClient

return mock.create_autospec(SpannerClient, instance=True)

def test_ctor_defaults(self):

from google.cloud.spanner_v1.pool import BurstyPool

instance = _Instance(self.INSTANCE_NAME)

database = self._make_one(self.DATABASE_ID, instance)

self.assertEqual(database.database_id, self.DATABASE_ID)

self.assertIs(database._instance, instance)

self.assertEqual(list(database.ddl_statements), [])

self.assertIsInstance(database._pool, BurstyPool)

# BurstyPool does not create sessions during 'bind()'.

self.assertTrue(database._pool._sessions.empty())

def test_ctor_w_explicit_pool(self):

instance = _Instance(self.INSTANCE_NAME)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

self.assertEqual(database.database_id, self.DATABASE_ID)

self.assertIs(database._instance, instance)

self.assertEqual(list(database.ddl_statements), [])

self.assertIs(database._pool, pool)

self.assertIs(pool._bound, database)

def test_ctor_w_ddl_statements_non_string(self):

with self.assertRaises(ValueError):

self._make_one(

self.DATABASE_ID, instance=object(), ddl_statements=[object()]

)

def test_ctor_w_ddl_statements_w_create_database(self):

with self.assertRaises(ValueError):

self._make_one(

self.DATABASE_ID,

instance=object(),

ddl_statements=["CREATE DATABASE foo"],

)

def test_ctor_w_ddl_statements_ok(self):

from tests._fixtures import DDL_STATEMENTS

instance = _Instance(self.INSTANCE_NAME)

pool = _Pool()

database = self._make_one(

self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool

)

self.assertEqual(database.database_id, self.DATABASE_ID)

self.assertIs(database._instance, instance)

self.assertEqual(list(database.ddl_statements), DDL_STATEMENTS)

def test_from_pb_bad_database_name(self):

from google.cloud.spanner_admin_database_v1 import Database

database_name = "INCORRECT_FORMAT"

database_pb = Database(name=database_name)

klass = self._get_target_class()

with self.assertRaises(ValueError):

klass.from_pb(database_pb, None)

def test_from_pb_project_mistmatch(self):

from google.cloud.spanner_admin_database_v1 import Database

ALT_PROJECT = "ALT_PROJECT"

client = _Client(project=ALT_PROJECT)

instance = _Instance(self.INSTANCE_NAME, client)

database_pb = Database(name=self.DATABASE_NAME)

klass = self._get_target_class()

with self.assertRaises(ValueError):

klass.from_pb(database_pb, instance)

def test_from_pb_instance_mistmatch(self):

from google.cloud.spanner_admin_database_v1 import Database

ALT_INSTANCE = "/projects/%s/instances/ALT-INSTANCE" % (self.PROJECT_ID,)

client = _Client()

instance = _Instance(ALT_INSTANCE, client)

database_pb = Database(name=self.DATABASE_NAME)

klass = self._get_target_class()

with self.assertRaises(ValueError):

klass.from_pb(database_pb, instance)

def test_from_pb_success_w_explicit_pool(self):

from google.cloud.spanner_admin_database_v1 import Database

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client)

database_pb = Database(name=self.DATABASE_NAME)

klass = self._get_target_class()

pool = _Pool()

database = klass.from_pb(database_pb, instance, pool=pool)

self.assertIsInstance(database, klass)

self.assertEqual(database._instance, instance)

self.assertEqual(database.database_id, self.DATABASE_ID)

self.assertIs(database._pool, pool)

def test_from_pb_success_w_hyphen_w_default_pool(self):

from google.cloud.spanner_admin_database_v1 import Database

from google.cloud.spanner_v1.pool import BurstyPool

DATABASE_ID_HYPHEN = "database-id"

DATABASE_NAME_HYPHEN = self.INSTANCE_NAME + "/databases/" + DATABASE_ID_HYPHEN

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client)

database_pb = Database(name=DATABASE_NAME_HYPHEN)

klass = self._get_target_class()

database = klass.from_pb(database_pb, instance)

self.assertIsInstance(database, klass)

self.assertEqual(database._instance, instance)

self.assertEqual(database.database_id, DATABASE_ID_HYPHEN)

self.assertIsInstance(database._pool, BurstyPool)

# BurstyPool does not create sessions during 'bind()'.

self.assertTrue(database._pool._sessions.empty())

def test_name_property(self):

instance = _Instance(self.INSTANCE_NAME)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

expected_name = self.DATABASE_NAME

self.assertEqual(database.name, expected_name)

def test_create_time_property(self):

instance = _Instance(self.INSTANCE_NAME)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

expected_create_time = database._create_time = self._make_timestamp()

self.assertEqual(database.create_time, expected_create_time)

def test_state_property(self):

from google.cloud.spanner_admin_database_v1 import Database

instance = _Instance(self.INSTANCE_NAME)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

expected_state = database._state = Database.State.READY

self.assertEqual(database.state, expected_state)

def test_restore_info(self):

from google.cloud.spanner_admin_database_v1 import RestoreInfo

instance = _Instance(self.INSTANCE_NAME)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

restore_info = database._restore_info = mock.create_autospec(

RestoreInfo, instance=True

)

self.assertEqual(database.restore_info, restore_info)

def test_spanner_api_property_w_scopeless_creds(self):

client = _Client()

client_info = client._client_info = mock.Mock()

client_options = client._client_options = mock.Mock()

credentials = client.credentials = object()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

patch = mock.patch("google.cloud.spanner_v1.database.SpannerClient")

with patch as spanner_client:

api = database.spanner_api

self.assertIs(api, spanner_client.return_value)

# API instance is cached

again = database.spanner_api

self.assertIs(again, api)

spanner_client.assert_called_once_with(

credentials=credentials,

client_info=client_info,

client_options=client_options,

)

def test_spanner_api_w_scoped_creds(self):

import google.auth.credentials

from google.cloud.spanner_v1.database import SPANNER_DATA_SCOPE

class _CredentialsWithScopes(google.auth.credentials.Scoped):

def __init__(self, scopes=(), source=None):

self._scopes = scopes

self._source = source

def requires_scopes(self): # pragma: NO COVER

return True

def with_scopes(self, scopes):

return self.__class__(scopes, self)

expected_scopes = (SPANNER_DATA_SCOPE,)

client = _Client()

client_info = client._client_info = mock.Mock()

client_options = client._client_options = mock.Mock()

credentials = client.credentials = _CredentialsWithScopes()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

patch = mock.patch("google.cloud.spanner_v1.database.SpannerClient")

with patch as spanner_client:

api = database.spanner_api

# API instance is cached

again = database.spanner_api

self.assertIs(again, api)

self.assertEqual(len(spanner_client.call_args_list), 1)

called_args, called_kw = spanner_client.call_args

self.assertEqual(called_args, ())

self.assertEqual(called_kw["client_info"], client_info)

self.assertEqual(called_kw["client_options"], client_options)

scoped = called_kw["credentials"]

self.assertEqual(scoped._scopes, expected_scopes)

self.assertIs(scoped._source, credentials)

def test_spanner_api_w_emulator_host(self):

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client, emulator_host="host")

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

patch = mock.patch("google.cloud.spanner_v1.database.SpannerClient")

with patch as spanner_client:

api = database.spanner_api

self.assertIs(api, spanner_client.return_value)

# API instance is cached

again = database.spanner_api

self.assertIs(again, api)

self.assertEqual(len(spanner_client.call_args_list), 1)

called_args, called_kw = spanner_client.call_args

self.assertEqual(called_args, ())

self.assertIsNotNone(called_kw["transport"])

def test___eq__(self):

instance = _Instance(self.INSTANCE_NAME)

pool1, pool2 = _Pool(), _Pool()

database1 = self._make_one(self.DATABASE_ID, instance, pool=pool1)

database2 = self._make_one(self.DATABASE_ID, instance, pool=pool2)

self.assertEqual(database1, database2)

def test___eq__type_differ(self):

pool = _Pool()

database1 = self._make_one(self.DATABASE_ID, None, pool=pool)

database2 = object()

self.assertNotEqual(database1, database2)

def test___ne__same_value(self):

instance = _Instance(self.INSTANCE_NAME)

pool1, pool2 = _Pool(), _Pool()

database1 = self._make_one(self.DATABASE_ID, instance, pool=pool1)

database2 = self._make_one(self.DATABASE_ID, instance, pool=pool2)

comparison_val = database1 != database2

self.assertFalse(comparison_val)

def test___ne__(self):

pool1, pool2 = _Pool(), _Pool()

database1 = self._make_one("database_id1", "instance1", pool=pool1)

database2 = self._make_one("database_id2", "instance2", pool=pool2)

self.assertNotEqual(database1, database2)

def test_create_grpc_error(self):

from google.api_core.exceptions import GoogleAPICallError

from google.api_core.exceptions import Unknown

from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.create_database.side_effect = Unknown("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(GoogleAPICallError):

database.create()

expected_request = CreateDatabaseRequest(

parent=self.INSTANCE_NAME,

create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),

extra_statements=[],

)

api.create_database.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_create_already_exists(self):

from google.cloud.exceptions import Conflict

from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest

DATABASE_ID_HYPHEN = "database-id"

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.create_database.side_effect = Conflict("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(DATABASE_ID_HYPHEN, instance, pool=pool)

with self.assertRaises(Conflict):

database.create()

expected_request = CreateDatabaseRequest(

parent=self.INSTANCE_NAME,

create_statement="CREATE DATABASE `{}`".format(DATABASE_ID_HYPHEN),

extra_statements=[],

)

api.create_database.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_create_instance_not_found(self):

from google.cloud.exceptions import NotFound

from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.create_database.side_effect = NotFound("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(NotFound):

database.create()

expected_request = CreateDatabaseRequest(

parent=self.INSTANCE_NAME,

create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),

extra_statements=[],

)

api.create_database.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_create_success(self):

from tests._fixtures import DDL_STATEMENTS

from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest

op_future = object()

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.create_database.return_value = op_future

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(

self.DATABASE_ID, instance, ddl_statements=DDL_STATEMENTS, pool=pool

)

future = database.create()

self.assertIs(future, op_future)

expected_request = CreateDatabaseRequest(

parent=self.INSTANCE_NAME,

create_statement="CREATE DATABASE {}".format(self.DATABASE_ID),

extra_statements=DDL_STATEMENTS,

)

api.create_database.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_exists_grpc_error(self):

from google.api_core.exceptions import Unknown

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.get_database_ddl.side_effect = Unknown("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(Unknown):

database.exists()

api.get_database_ddl.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_exists_not_found(self):

from google.cloud.exceptions import NotFound

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.get_database_ddl.side_effect = NotFound("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

self.assertFalse(database.exists())

api.get_database_ddl.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_exists_success(self):

from google.cloud.spanner_admin_database_v1 import GetDatabaseDdlResponse

from tests._fixtures import DDL_STATEMENTS

client = _Client()

ddl_pb = GetDatabaseDdlResponse(statements=DDL_STATEMENTS)

api = client.database_admin_api = self._make_database_admin_api()

api.get_database_ddl.return_value = ddl_pb

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

self.assertTrue(database.exists())

api.get_database_ddl.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_reload_grpc_error(self):

from google.api_core.exceptions import Unknown

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.get_database_ddl.side_effect = Unknown("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(Unknown):

database.reload()

api.get_database_ddl.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_reload_not_found(self):

from google.cloud.exceptions import NotFound

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.get_database_ddl.side_effect = NotFound("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(NotFound):

database.reload()

api.get_database_ddl.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_reload_success(self):

from google.cloud.spanner_admin_database_v1 import Database

from google.cloud.spanner_admin_database_v1 import GetDatabaseDdlResponse

from google.cloud.spanner_admin_database_v1 import RestoreInfo

from google.cloud._helpers import _datetime_to_pb_timestamp

from tests._fixtures import DDL_STATEMENTS

timestamp = self._make_timestamp()

restore_info = RestoreInfo()

client = _Client()

ddl_pb = GetDatabaseDdlResponse(statements=DDL_STATEMENTS)

api = client.database_admin_api = self._make_database_admin_api()

api.get_database_ddl.return_value = ddl_pb

db_pb = Database(

state=2,

create_time=_datetime_to_pb_timestamp(timestamp),

restore_info=restore_info,

)

api.get_database.return_value = db_pb

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

database.reload()

self.assertEqual(database._state, Database.State.READY)

self.assertEqual(database._create_time, timestamp)

self.assertEqual(database._restore_info, restore_info)

self.assertEqual(database._ddl_statements, tuple(DDL_STATEMENTS))

api.get_database_ddl.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

api.get_database.assert_called_once_with(

name=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_update_ddl_grpc_error(self):

from google.api_core.exceptions import Unknown

from tests._fixtures import DDL_STATEMENTS

from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.update_database_ddl.side_effect = Unknown("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(Unknown):

database.update_ddl(DDL_STATEMENTS)

expected_request = UpdateDatabaseDdlRequest(

database=self.DATABASE_NAME, statements=DDL_STATEMENTS, operation_id="",

)

api.update_database_ddl.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_update_ddl_not_found(self):

from google.cloud.exceptions import NotFound

from tests._fixtures import DDL_STATEMENTS

from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.update_database_ddl.side_effect = NotFound("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(NotFound):

database.update_ddl(DDL_STATEMENTS)

expected_request = UpdateDatabaseDdlRequest(

database=self.DATABASE_NAME, statements=DDL_STATEMENTS, operation_id="",

)

api.update_database_ddl.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_update_ddl(self):

from tests._fixtures import DDL_STATEMENTS

from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest

op_future = object()

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.update_database_ddl.return_value = op_future

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

future = database.update_ddl(DDL_STATEMENTS)

self.assertIs(future, op_future)

expected_request = UpdateDatabaseDdlRequest(

database=self.DATABASE_NAME, statements=DDL_STATEMENTS, operation_id="",

)

api.update_database_ddl.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_update_ddl_w_operation_id(self):

from tests._fixtures import DDL_STATEMENTS

from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest

op_future = object()

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.update_database_ddl.return_value = op_future

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

future = database.update_ddl(DDL_STATEMENTS, operation_id="someOperationId")

self.assertIs(future, op_future)

expected_request = UpdateDatabaseDdlRequest(

database=self.DATABASE_NAME,

statements=DDL_STATEMENTS,

operation_id="someOperationId",

)

api.update_database_ddl.assert_called_once_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_drop_grpc_error(self):

from google.api_core.exceptions import Unknown

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.drop_database.side_effect = Unknown("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(Unknown):

database.drop()

api.drop_database.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_drop_not_found(self):

from google.cloud.exceptions import NotFound

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.drop_database.side_effect = NotFound("testing")

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

with self.assertRaises(NotFound):

database.drop()

api.drop_database.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def test_drop_success(self):

from google.protobuf.empty_pb2 import Empty

client = _Client()

api = client.database_admin_api = self._make_database_admin_api()

api.drop_database.return_value = Empty()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

database.drop()

api.drop_database.assert_called_once_with(

database=self.DATABASE_NAME,

metadata=[("google-cloud-resource-prefix", database.name)],

)

def _execute_partitioned_dml_helper(

self, dml, params=None, param_types=None, query_options=None, retried=False

):

from google.api_core.exceptions import Aborted

from google.api_core.retry import Retry

from google.protobuf.struct_pb2 import Struct

from google.cloud.spanner_v1 import (

PartialResultSet,

ResultSetStats,

)

from google.cloud.spanner_v1 import (

Transaction as TransactionPB,

TransactionSelector,

TransactionOptions,

)

from google.cloud.spanner_v1._helpers import (

_make_value_pb,

_merge_query_options,

)

from google.cloud.spanner_v1 import ExecuteSqlRequest

import collections

MethodConfig = collections.namedtuple("MethodConfig", ["retry"])

transaction_pb = TransactionPB(id=self.TRANSACTION_ID)

stats_pb = ResultSetStats(row_count_lower_bound=2)

result_sets = [PartialResultSet(stats=stats_pb)]

iterator = _MockIterator(*result_sets)

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

session = _Session()

pool.put(session)

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

api = database._spanner_api = self._make_spanner_api()

api._method_configs = {"ExecuteStreamingSql": MethodConfig(retry=Retry())}

if retried:

retry_transaction_pb = TransactionPB(id=self.RETRY_TRANSACTION_ID)

api.begin_transaction.side_effect = [transaction_pb, retry_transaction_pb]

api.execute_streaming_sql.side_effect = [Aborted("test"), iterator]

else:

api.begin_transaction.return_value = transaction_pb

api.execute_streaming_sql.return_value = iterator

row_count = database.execute_partitioned_dml(

dml, params, param_types, query_options

)

self.assertEqual(row_count, 2)

txn_options = TransactionOptions(

partitioned_dml=TransactionOptions.PartitionedDml()

)

api.begin_transaction.assert_called_with(

session=session.name,

options=txn_options,

metadata=[("google-cloud-resource-prefix", database.name)],

)

if retried:

self.assertEqual(api.begin_transaction.call_count, 2)

else:

self.assertEqual(api.begin_transaction.call_count, 1)

if params:

expected_params = Struct(

fields={key: _make_value_pb(value) for (key, value) in params.items()}

)

else:

expected_params = {}

expected_transaction = TransactionSelector(id=self.TRANSACTION_ID)

expected_query_options = client._query_options

if query_options:

expected_query_options = _merge_query_options(

expected_query_options, query_options

)

expected_request = ExecuteSqlRequest(

session=self.SESSION_NAME,

sql=dml,

transaction=expected_transaction,

params=expected_params,

param_types=param_types,

query_options=expected_query_options,

)

api.execute_streaming_sql.assert_any_call(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

if retried:

expected_retry_transaction = TransactionSelector(

id=self.RETRY_TRANSACTION_ID

)

expected_request = ExecuteSqlRequest(

session=self.SESSION_NAME,

sql=dml,

transaction=expected_retry_transaction,

params=expected_params,

param_types=param_types,

query_options=expected_query_options,

)

api.execute_streaming_sql.assert_called_with(

request=expected_request,

metadata=[("google-cloud-resource-prefix", database.name)],

)

self.assertEqual(api.execute_streaming_sql.call_count, 2)

else:

self.assertEqual(api.execute_streaming_sql.call_count, 1)

def test_execute_partitioned_dml_wo_params(self):

self._execute_partitioned_dml_helper(dml=DML_WO_PARAM)

def test_execute_partitioned_dml_w_params_wo_param_types(self):

with self.assertRaises(ValueError):

self._execute_partitioned_dml_helper(dml=DML_W_PARAM, params=PARAMS)

def test_execute_partitioned_dml_w_params_and_param_types(self):

self._execute_partitioned_dml_helper(

dml=DML_W_PARAM, params=PARAMS, param_types=PARAM_TYPES

)

def test_execute_partitioned_dml_w_query_options(self):

from google.cloud.spanner_v1 import ExecuteSqlRequest

self._execute_partitioned_dml_helper(

dml=DML_W_PARAM,

query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="3"),

)

def test_execute_partitioned_dml_wo_params_retry_aborted(self):

self._execute_partitioned_dml_helper(dml=DML_WO_PARAM, retried=True)

def test_session_factory_defaults(self):

from google.cloud.spanner_v1.session import Session

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

session = database.session()

self.assertIsInstance(session, Session)

self.assertIs(session.session_id, None)

self.assertIs(session._database, database)

self.assertEqual(session.labels, {})

def test_session_factory_w_labels(self):

from google.cloud.spanner_v1.session import Session

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

labels = {"foo": "bar"}

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

session = database.session(labels=labels)

self.assertIsInstance(session, Session)

self.assertIs(session.session_id, None)

self.assertIs(session._database, database)

self.assertEqual(session.labels, labels)

def test_snapshot_defaults(self):

from google.cloud.spanner_v1.database import SnapshotCheckout

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

session = _Session()

pool.put(session)

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

checkout = database.snapshot()

self.assertIsInstance(checkout, SnapshotCheckout)

self.assertIs(checkout._database, database)

self.assertEqual(checkout._kw, {})

def test_snapshot_w_read_timestamp_and_multi_use(self):

import datetime

from google.cloud._helpers import UTC

from google.cloud.spanner_v1.database import SnapshotCheckout

now = datetime.datetime.utcnow().replace(tzinfo=UTC)

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

session = _Session()

pool.put(session)

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

checkout = database.snapshot(read_timestamp=now, multi_use=True)

self.assertIsInstance(checkout, SnapshotCheckout)

self.assertIs(checkout._database, database)

self.assertEqual(checkout._kw, {"read_timestamp": now, "multi_use": True})

def test_batch(self):

from google.cloud.spanner_v1.database import BatchCheckout

client = _Client()

instance = _Instance(self.INSTANCE_NAME, client=client)

pool = _Pool()

session = _Session()

pool.put(session)

database = self._make_one(self.DATABASE_ID, instance, pool=pool)

checkout = database.batch()

self.assertIsInstance(checkout, BatchCheckout)

self.assertIs(checkout._database, database)

def test_batch_snapshot(self):

from google.cloud.spanner_v1.database import BatchSnapshot

database = self._make_one(self.DATABASE_ID, instance=object(), pool=_Pool())

batch_txn = database.batch_snapshot()

self.assertIsInstance(batch_txn, BatchSnapshot)

self.assertIs(batch_txn._database, database)

self.assertIsNone(batch_txn._read_timestamp)

self.assertIsNone(batch_txn._exact_staleness)

def test_batch_snapshot_w_read_timestamp(self):

from google.cloud.spanner_v1.database import BatchSnapshot

database = self._make_one(self.DATABASE_ID, instance=object(), pool=_Pool())

timestamp = self._make_timestamp()

batch_txn = database.batch_snapshot(read_timestamp=timestamp)

self.assertIsInstance(batch_txn, BatchSnapshot)

self.assertIs(batch_txn._database, database)

self.assertEqual(batch_txn._read_timestamp, timestamp)

self.assertIsNone(batch_txn._exact_staleness)

def test_batch_snapshot_w_exact_staleness(self):

from google.cloud.spanner_v1.database import BatchSnapshot

database = self._make_one(self.DATABASE_ID, instance=object(), pool=_Pool())