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
