def test_connect_instance_not_found(self): from google.cloud.spanner_dbapi import connect with mock.patch("google.cloud.spanner_v1.instance.Instance.exists", return_value=False): with self.assertRaises(ValueError): connect("test-instance", "test-database")
def test_instance_not_found(self): with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=False, ) as exists_mock: with self.assertRaises(ValueError): connect("test-instance", "test-database") exists_mock.assert_called_once_with()
def test_w_instance_not_found(self, mock_client): from google.cloud.spanner_dbapi import connect client = mock_client.return_value instance = client.instance.return_value instance.exists.return_value = False with self.assertRaises(ValueError): connect(INSTANCE, DATABASE) instance.exists.assert_called_once_with()
def test_sessions_pool(self): database_id = "test-database" pool = FixedSizePool() with mock.patch("google.cloud.spanner_v1.instance.Instance.database" ) as database_mock: with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True, ): connect("test-instance", database_id, pool=pool) database_mock.assert_called_once_with(database_id, pool=pool)
def test_executemany(self): from google.cloud.spanner_dbapi import connect operation = """SELECT * FROM table1 WHERE "col1" = @a1""" params_seq = ((1,), (2,)) with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True ): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True ): connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._result_set = [1, 2, 3] cursor._itr = iter([1, 2, 3]) with mock.patch( "google.cloud.spanner_dbapi.cursor.Cursor.execute" ) as execute_mock: cursor.executemany(operation, params_seq) execute_mock.assert_has_calls( (mock.call(operation, (1,)), mock.call(operation, (2,))) )
def test_retry_aborted_retry_without_delay(self, mock_client): """ Check that in case of a retried transaction failed, the connection will retry it once again. """ from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.cursor import Statement row = ["field1", "field2"] connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._checksum = ResultsChecksum() cursor._checksum.consume_result(row) statement = Statement("SELECT 1", [], {}, cursor._checksum, False) connection._statements.append(statement) metadata_mock = mock.Mock() metadata_mock.trailing_metadata.return_value = {} run_mock = connection.run_statement = mock.Mock() run_mock.side_effect = [ Aborted("Aborted", errors=[metadata_mock]), ([row], ResultsChecksum()), ] connection._get_retry_delay = mock.Mock(return_value=False) connection.retry_transaction() run_mock.assert_has_calls( (mock.call(statement, retried=True), mock.call(statement, retried=True),) )
def test_commit_retry_aborted_statements(self, mock_client): """Check that retried transaction executing the same statements.""" from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.cursor import Statement row = ["field1", "field2"] connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._checksum = ResultsChecksum() cursor._checksum.consume_result(row) statement = Statement("SELECT 1", [], {}, cursor._checksum, False) connection._statements.append(statement) mock_transaction = mock.Mock(rolled_back=False, committed=False) connection._transaction = mock_transaction mock_transaction.commit.side_effect = [Aborted("Aborted"), None] run_mock = connection.run_statement = mock.Mock() run_mock.return_value = ([row], ResultsChecksum()) connection.commit() run_mock.assert_called_with(statement, retried=True)
def test_close(self): from google.cloud.spanner_dbapi import connect, InterfaceError with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True, ): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True, ): connection = connect("test-instance", "test-database") self.assertFalse(connection.is_closed) connection.close() self.assertTrue(connection.is_closed) with self.assertRaises(InterfaceError): connection.cursor() connection._transaction = mock_transaction = mock.MagicMock() mock_transaction.committed = mock_transaction.rolled_back = False mock_transaction.rollback = mock_rollback = mock.MagicMock() connection.close() mock_rollback.assert_called_once_with()
def test_commit_retry_aborted_statements(self): """Check that retried transaction executing the same statements.""" from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.cursor import Statement row = ["field1", "field2"] with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True, ): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True, ): connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._checksum = ResultsChecksum() cursor._checksum.consume_result(row) statement = Statement("SELECT 1", [], {}, cursor._checksum, False) connection._statements.append(statement) connection._transaction = mock.Mock(rolled_back=False, committed=False) with mock.patch.object( connection._transaction, "commit", side_effect=(Aborted("Aborted"), None), ): with mock.patch( "google.cloud.spanner_dbapi.connection.Connection.run_statement", return_value=([row], ResultsChecksum()), ) as run_mock: connection.commit() run_mock.assert_called_with(statement, retried=True)
def test_fetchmany_retry_aborted(self): """Check that aborted fetch re-executing transaction.""" from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.connection import connect with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True, ): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True, ): connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._checksum = ResultsChecksum() with mock.patch( "google.cloud.spanner_dbapi.cursor.Cursor.__next__", side_effect=(Aborted("Aborted"), None), ): with mock.patch( "google.cloud.spanner_dbapi.connection.Connection.retry_transaction" ) as retry_mock: cursor.fetchmany() retry_mock.assert_called_with()
def test_executemany_insert_batch_non_autocommit(self): from google.cloud.spanner_dbapi import connect from google.cloud.spanner_v1.param_types import INT64 from google.cloud.spanner_v1.types.spanner import Session sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)""" with mock.patch("google.cloud.spanner_v1.instance.Instance.exists", return_value=True): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True, ): connection = connect("test-instance", "test-database") transaction = self._transaction_mock() cursor = connection.cursor() with mock.patch( "google.cloud.spanner_v1.services.spanner.client.SpannerClient.create_session", return_value=Session(), ): with mock.patch( "google.cloud.spanner_v1.session.Session.transaction", return_value=transaction, ): cursor.executemany(sql, [(1, 2, 3, 4), (5, 6, 7, 8)]) transaction.batch_update.assert_called_once_with([ ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", { "a0": 1, "a1": 2, "a2": 3, "a3": 4 }, { "a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64 }, ), ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", { "a0": 5, "a1": 6, "a2": 7, "a3": 8 }, { "a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64 }, ), ])
def test_ddls_with_semicolon(self, mock_client): """ Check that one script with several DDL statements separated with semicolons is splitted into several DDLs. """ from google.cloud.spanner_dbapi.connection import connect EXP_DDLS = [ "CREATE TABLE table_name (row_id INT64) PRIMARY KEY ()", "DROP INDEX index_name", ("CREATE TABLE papers (" "\n id INT64," "\n authors ARRAY<STRING(100)>," '\n author_list STRING(MAX) AS (ARRAY_TO_STRING(authors, ";")) stored' ") PRIMARY KEY (id)"), "DROP TABLE table_name", ] connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor.execute( "CREATE TABLE table_name (row_id INT64) PRIMARY KEY ();" "DROP INDEX index_name;\n" "CREATE TABLE papers (" "\n id INT64," "\n authors ARRAY<STRING(100)>," '\n author_list STRING(MAX) AS (ARRAY_TO_STRING(authors, ";")) stored' ") PRIMARY KEY (id);" "DROP TABLE table_name;", ) self.assertEqual(connection._ddl_statements, EXP_DDLS)
def test_peek_iterator_aborted_autocommit(self, mock_client): """ Checking that an Aborted exception is retried in case it happened while streaming the first element with a PeekIterator in autocommit mode. """ from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.connection import connect connection = connect("test-instance", "test-database") connection.autocommit = True cursor = connection.cursor() with mock.patch( "google.cloud.spanner_dbapi.utils.PeekIterator.__init__", side_effect=(Aborted("Aborted"), None), ): with mock.patch( "google.cloud.spanner_dbapi.connection.Connection.retry_transaction" ) as retry_mock: with mock.patch( "google.cloud.spanner_dbapi.connection.Connection.run_statement", return_value=((1, 2, 3), None), ): with mock.patch( "google.cloud.spanner_v1.database.Database.snapshot" ): cursor.execute("SELECT * FROM table_name") retry_mock.assert_called_with()
def test_w_explicit(self, mock_client): from google.cloud.spanner_v1.pool import AbstractSessionPool from google.cloud.spanner_dbapi import connect from google.cloud.spanner_dbapi import Connection from google.cloud.spanner_dbapi.version import PY_VERSION credentials = _make_credentials() pool = mock.create_autospec(AbstractSessionPool) client = mock_client.return_value instance = client.instance.return_value database = instance.database.return_value connection = connect( INSTANCE, DATABASE, PROJECT, credentials, pool=pool, user_agent=USER_AGENT, ) self.assertIsInstance(connection, Connection) mock_client.assert_called_once_with(project=PROJECT, credentials=credentials, client_info=mock.ANY) client_info = mock_client.call_args_list[0][1]["client_info"] self.assertEqual(client_info.user_agent, USER_AGENT) self.assertEqual(client_info.python_version, PY_VERSION) self.assertIs(connection.instance, instance) client.instance.assert_called_once_with(INSTANCE) self.assertIs(connection.database, database) instance.database.assert_called_once_with(DATABASE, pool=pool)
def test_fetchmany_retry_aborted_statements_checksums_mismatch( self, mock_client): """Check transaction retrying with underlying data being changed.""" from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.exceptions import RetryAborted from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.cursor import Statement row = ["field1", "field2"] row2 = ["updated_field1", "field2"] connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._checksum = ResultsChecksum() cursor._checksum.consume_result(row) statement = Statement("SELECT 1", [], {}, cursor._checksum, False) connection._statements.append(statement) with mock.patch( "google.cloud.spanner_dbapi.cursor.Cursor.__next__", side_effect=(Aborted("Aborted"), None), ): with mock.patch( "google.cloud.spanner_dbapi.connection.Connection.run_statement", return_value=([row2], ResultsChecksum()), ) as run_mock: with self.assertRaises(RetryAborted): cursor.fetchmany(len(row)) run_mock.assert_called_with(statement, retried=True)
def test_executemany_delete_batch_autocommit(self): from google.cloud.spanner_dbapi import connect from google.cloud.spanner_v1.param_types import INT64 from google.cloud.spanner_v1.types.spanner import Session sql = "DELETE FROM table WHERE col1 = %s" connection = connect("test-instance", "test-database") connection.autocommit = True transaction = self._transaction_mock() cursor = connection.cursor() with mock.patch( "google.cloud.spanner_v1.services.spanner.client.SpannerClient.create_session", return_value=Session(), ): with mock.patch( "google.cloud.spanner_v1.session.Session.transaction", return_value=transaction, ): cursor.executemany(sql, [(1,), (2,), (3,)]) transaction.batch_update.assert_called_once_with( [ ("DELETE FROM table WHERE col1 = @a0", {"a0": 1}, {"a0": INT64}), ("DELETE FROM table WHERE col1 = @a0", {"a0": 2}, {"a0": INT64}), ("DELETE FROM table WHERE col1 = @a0", {"a0": 3}, {"a0": INT64}), ] )
def enable_autocommit_mode(instance_id, database_id): """Enables autocommit mode.""" # [START spanner_enable_autocommit_mode] connection = connect(instance_id, database_id) connection.autocommit = True print("Autocommit mode is enabled.") cursor = connection.cursor() cursor.execute("""CREATE TABLE Singers ( SingerId INT64 NOT NULL, FirstName STRING(1024), LastName STRING(1024), SingerInfo BYTES(MAX) ) PRIMARY KEY (SingerId)""") cursor.execute( """INSERT INTO Singers (SingerId, FirstName, LastName) VALUES (12, 'Melissa', 'Garcia'), (13, 'Russell', 'Morales'), (14, 'Jacqueline', 'Long'), (15, 'Dylan', 'Shaw')""") cursor.execute("""SELECT * FROM Singers WHERE SingerId = 13""") print( "SingerId: {}, AlbumId: {}, AlbumTitle: {}".format(*cursor.fetchone())) connection.close()
def test_ddls_with_semicolon(self): """ Check that one script with several DDL statements separated with semicolons is splitted into several DDLs. """ from google.cloud.spanner_dbapi.connection import connect EXP_DDLS = [ "CREATE TABLE table_name (row_id INT64) PRIMARY KEY ()", "DROP INDEX index_name", "DROP TABLE table_name", ] with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True, ): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True, ): connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor.execute("CREATE TABLE table_name (row_id INT64) PRIMARY KEY ();" "DROP INDEX index_name;\n" "DROP TABLE table_name;") self.assertEqual(connection._ddl_statements, EXP_DDLS)
def test_connect(self): PROJECT = "test-project" USER_AGENT = "user-agent" CREDENTIALS = _make_credentials() CLIENT_INFO = ClientInfo(user_agent=USER_AGENT) with mock.patch( "google.cloud.spanner_dbapi.spanner_v1.Client") as client_mock: with mock.patch( "google.cloud.spanner_dbapi.google_client_info", return_value=CLIENT_INFO, ) as client_info_mock: connection = connect( "test-instance", "test-database", PROJECT, CREDENTIALS, USER_AGENT, ) self.assertIsInstance(connection, Connection) client_info_mock.assert_called_once_with(USER_AGENT) client_mock.assert_called_once_with( project=PROJECT, credentials=CREDENTIALS, client_info=CLIENT_INFO, )
def test_close(self, mock_client): from google.cloud.spanner_dbapi import connect from google.cloud.spanner_dbapi import InterfaceError connection = connect("test-instance", "test-database") self.assertFalse(connection.is_closed) connection.close() self.assertTrue(connection.is_closed) with self.assertRaises(InterfaceError): connection.cursor() mock_transaction = mock.MagicMock() mock_transaction.committed = mock_transaction.rolled_back = False connection._transaction = mock_transaction mock_rollback = mock.MagicMock() mock_transaction.rollback = mock_rollback connection.close() mock_rollback.assert_called_once_with() connection._transaction = mock.MagicMock() connection._own_pool = False connection.close() self.assertTrue(connection.is_closed)
def test_w_credential_file_path(self, mock_client): from google.cloud.spanner_dbapi import connect from google.cloud.spanner_dbapi import Connection from google.cloud.spanner_dbapi.version import PY_VERSION credentials_path = "dummy/file/path.json" connection = connect( INSTANCE, DATABASE, PROJECT, credentials=credentials_path, user_agent=USER_AGENT, ) self.assertIsInstance(connection, Connection) factory = mock_client.from_service_account_json factory.assert_called_once_with( credentials_path, project=PROJECT, client_info=mock.ANY, ) client_info = factory.call_args_list[0][1]["client_info"] self.assertEqual(client_info.user_agent, USER_AGENT) self.assertEqual(client_info.python_version, PY_VERSION)
def test_executemany_insert_batch_failed(self): from google.cloud.spanner_dbapi import connect from google.cloud.spanner_dbapi.exceptions import OperationalError from google.cloud.spanner_v1.types.spanner import Session from google.rpc.code_pb2 import UNKNOWN sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)""" err_details = "Details here" with mock.patch("google.cloud.spanner_v1.instance.Instance.exists", return_value=True): with mock.patch( "google.cloud.spanner_v1.database.Database.exists", return_value=True, ): connection = connect("test-instance", "test-database") connection.autocommit = True cursor = connection.cursor() transaction = mock.Mock(committed=False, rolled_back=False) transaction.batch_update = mock.Mock( return_value=(mock.Mock(code=UNKNOWN, details=err_details), [])) with mock.patch( "google.cloud.spanner_v1.services.spanner.client.SpannerClient.create_session", return_value=Session(), ): with mock.patch( "google.cloud.spanner_v1.session.Session.transaction", return_value=transaction, ): with self.assertRaisesRegex(OperationalError, err_details): cursor.executemany(sql, [(1, 2, 3, 4), (5, 6, 7, 8)])
def test_fetchall_retry_aborted_statements(self, mock_client): """Check that retried transaction executing the same statements.""" from google.api_core.exceptions import Aborted from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.cursor import Statement row = ["field1", "field2"] connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor._checksum = ResultsChecksum() cursor._checksum.consume_result(row) statement = Statement("SELECT 1", [], {}, cursor._checksum, False) connection._statements.append(statement) with mock.patch( "google.cloud.spanner_dbapi.cursor.Cursor.__iter__", side_effect=(Aborted("Aborted"), iter(row)), ): with mock.patch( "google.cloud.spanner_dbapi.connection.Connection.run_statement", return_value=([row], ResultsChecksum()), ) as run_mock: cursor.fetchall() run_mock.assert_called_with(statement, retried=True)
def _cleanup(self): """Drop the test table.""" conn = spanner_dbapi.connect(INSTANCE, DATABASE) try: conn.database.update_ddl(["DROP TABLE Singers"]) except NotFound: pass conn.close()
def test_executemany_update_batch_autocommit(self): from google.cloud.spanner_dbapi import connect from google.cloud.spanner_v1.param_types import INT64, STRING from google.cloud.spanner_v1.types.spanner import Session sql = "UPDATE table SET col1 = %s WHERE col2 = %s" connection = connect("test-instance", "test-database") connection.autocommit = True transaction = self._transaction_mock() cursor = connection.cursor() with mock.patch( "google.cloud.spanner_v1.services.spanner.client.SpannerClient.create_session", return_value=Session(), ): with mock.patch( "google.cloud.spanner_v1.session.Session.transaction", return_value=transaction, ): cursor.executemany(sql, [(1, "a"), (2, "b"), (3, "c")]) transaction.batch_update.assert_called_once_with([ ( "UPDATE table SET col1 = @a0 WHERE col2 = @a1", { "a0": 1, "a1": "a" }, { "a0": INT64, "a1": STRING }, ), ( "UPDATE table SET col1 = @a0 WHERE col2 = @a1", { "a0": 2, "a1": "b" }, { "a0": INT64, "a1": STRING }, ), ( "UPDATE table SET col1 = @a0 WHERE col2 = @a1", { "a0": 3, "a1": "c" }, { "a0": INT64, "a1": STRING }, ), ])
def test_default_sessions_pool(self): with mock.patch("google.cloud.spanner_v1.instance.Instance.database"): with mock.patch( "google.cloud.spanner_v1.instance.Instance.exists", return_value=True, ): connection = connect("test-instance", "test-database") self.assertIsNotNone(connection.database._pool)
def test_executemany_DLL(self, mock_client): from google.cloud.spanner_dbapi import connect, ProgrammingError connection = connect("test-instance", "test-database") cursor = connection.cursor() with self.assertRaises(ProgrammingError): cursor.executemany("""DROP DATABASE database_name""", ())
def test_connect_instance_id(self): INSTANCE = "test-instance" with mock.patch("google.cloud.spanner_v1.client.Client.instance" ) as instance_mock: connection = connect(INSTANCE, "test-database") instance_mock.assert_called_once_with(INSTANCE) self.assertIsInstance(connection, Connection)
def test_executemany_on_closed_cursor(self, mock_client): from google.cloud.spanner_dbapi import InterfaceError from google.cloud.spanner_dbapi import connect connection = connect("test-instance", "test-database") cursor = connection.cursor() cursor.close() with self.assertRaises(InterfaceError): cursor.executemany("""SELECT * FROM table1 WHERE "col1" = @a1""", ())
def test_enable_autocommit_mode(capsys, database): connection = connect(INSTANCE_ID, DATABASE_ID) cursor = connection.cursor() with mock.patch( "google.cloud.spanner_dbapi.connection.Cursor", return_value=cursor, ): autocommit.enable_autocommit_mode(INSTANCE_ID, DATABASE_ID) out, _ = capsys.readouterr() assert "Autocommit mode is enabled." in out assert "SingerId: 13, AlbumId: Russell, AlbumTitle: Morales" in out