test_cursor.py

# Copyright 2020 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.

"""Cursor() class unit tests."""

import mock
import sys
import unittest


class TestCursor(unittest.TestCase):

    INSTANCE = "test-instance"
    DATABASE = "test-database"

    def _get_target_class(self):
        from google.cloud.spanner_dbapi import Cursor

        return Cursor

    def _make_one(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    def _make_connection(self, *args, **kwargs):
        from google.cloud.spanner_dbapi import Connection

        return Connection(*args, **kwargs)

    def test_property_connection(self):
        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        self.assertEqual(cursor.connection, connection)

    def test_property_description(self):
        from google.cloud.spanner_dbapi._helpers import ColumnInfo

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)

        self.assertIsNone(cursor.description)
        cursor._result_set = res_set = mock.MagicMock()
        res_set.metadata.row_type.fields = [mock.MagicMock()]
        self.assertIsNotNone(cursor.description)
        self.assertIsInstance(cursor.description[0], ColumnInfo)

    def test_property_rowcount(self):
        from google.cloud.spanner_dbapi.cursor import _UNSET_COUNT

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        self.assertEqual(cursor.rowcount, _UNSET_COUNT)

    def test_callproc(self):
        from google.cloud.spanner_dbapi.exceptions import InterfaceError

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        cursor._is_closed = True
        with self.assertRaises(InterfaceError):
            cursor.callproc(procname=None)

    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(self.INSTANCE, self.DATABASE)

        cursor = connection.cursor()
        self.assertFalse(cursor.is_closed)

        cursor.close()

        self.assertTrue(cursor.is_closed)
        with self.assertRaises(InterfaceError):
            cursor.execute("SELECT * FROM database")

    def test_do_execute_update(self):
        from google.cloud.spanner_dbapi.cursor import _UNSET_COUNT

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        transaction = mock.MagicMock()

        def run_helper(ret_value):
            transaction.execute_update.return_value = ret_value
            res = cursor._do_execute_update(
                transaction=transaction, sql="SELECT * WHERE true", params={},
            )
            return res

        expected = "good"
        self.assertEqual(run_helper(expected), expected)
        self.assertEqual(cursor._row_count, _UNSET_COUNT)

        expected = 1234
        self.assertEqual(run_helper(expected), expected)
        self.assertEqual(cursor._row_count, expected)

    def test_execute_programming_error(self):
        from google.cloud.spanner_dbapi.exceptions import ProgrammingError

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        cursor.connection = None
        with self.assertRaises(ProgrammingError):
            cursor.execute(sql="")

    def test_execute_attribute_error(self):
        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)

        with self.assertRaises(AttributeError):
            cursor.execute(sql="SELECT 1")

    def test_execute_autocommit_off(self):
        from google.cloud.spanner_dbapi.utils import PeekIterator

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor.connection._autocommit = False
        cursor.connection.transaction_checkout = mock.MagicMock(autospec=True)

        cursor.execute("sql")
        self.assertIsInstance(cursor._result_set, mock.MagicMock)
        self.assertIsInstance(cursor._itr, PeekIterator)

    def test_execute_insert_statement_autocommit_off(self):
        from google.cloud.spanner_dbapi import parse_utils
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum
        from google.cloud.spanner_dbapi.utils import PeekIterator

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor.connection._autocommit = False
        cursor.connection.transaction_checkout = mock.MagicMock(autospec=True)

        cursor._checksum = ResultsChecksum()
        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            return_value=parse_utils.STMT_INSERT,
        ):
            with mock.patch(
                "google.cloud.spanner_dbapi.connection.Connection.run_statement",
                return_value=(mock.MagicMock(), ResultsChecksum()),
            ):
                cursor.execute(
                    sql="INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)"
                )
                self.assertIsInstance(cursor._result_set, mock.MagicMock)
                self.assertIsInstance(cursor._itr, PeekIterator)

    def test_execute_statement(self):
        from google.cloud.spanner_dbapi import parse_utils

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            side_effect=[parse_utils.STMT_DDL, parse_utils.STMT_INSERT],
        ) as mock_classify_stmt:
            sql = "sql"
            with self.assertRaises(ValueError):
                cursor.execute(sql=sql)
            mock_classify_stmt.assert_called_with(sql)
            self.assertEqual(mock_classify_stmt.call_count, 2)
            self.assertEqual(cursor.connection._ddl_statements, [])

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            return_value=parse_utils.STMT_DDL,
        ) as mock_classify_stmt:
            sql = "sql"
            cursor.execute(sql=sql)
            mock_classify_stmt.assert_called_with(sql)
            self.assertEqual(mock_classify_stmt.call_count, 2)
            self.assertEqual(cursor.connection._ddl_statements, [sql])

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            return_value=parse_utils.STMT_NON_UPDATING,
        ):
            with mock.patch(
                "google.cloud.spanner_dbapi.cursor.Cursor._handle_DQL",
                return_value=parse_utils.STMT_NON_UPDATING,
            ) as mock_handle_ddl:
                connection.autocommit = True
                sql = "sql"
                cursor.execute(sql=sql)
                mock_handle_ddl.assert_called_once_with(sql, None)

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            return_value=parse_utils.STMT_INSERT,
        ):
            with mock.patch(
                "google.cloud.spanner_dbapi._helpers.handle_insert",
                return_value=parse_utils.STMT_INSERT,
            ) as mock_handle_insert:
                sql = "sql"
                cursor.execute(sql=sql)
                mock_handle_insert.assert_called_once_with(connection, sql, None)

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            return_value="other_statement",
        ):
            cursor.connection._database = mock_db = mock.MagicMock()
            mock_db.run_in_transaction = mock_run_in = mock.MagicMock()
            sql = "sql"
            cursor.execute(sql=sql)
            mock_run_in.assert_called_once_with(cursor._do_execute_update, sql, None)

    def test_execute_integrity_error(self):
        from google.api_core import exceptions
        from google.cloud.spanner_dbapi.exceptions import IntegrityError

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            side_effect=exceptions.AlreadyExists("message"),
        ):
            with self.assertRaises(IntegrityError):
                cursor.execute(sql="sql")

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            side_effect=exceptions.FailedPrecondition("message"),
        ):
            with self.assertRaises(IntegrityError):
                cursor.execute(sql="sql")

    def test_execute_invalid_argument(self):
        from google.api_core import exceptions
        from google.cloud.spanner_dbapi.exceptions import ProgrammingError

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            side_effect=exceptions.InvalidArgument("message"),
        ):
            with self.assertRaises(ProgrammingError):
                cursor.execute(sql="sql")

    def test_execute_internal_server_error(self):
        from google.api_core import exceptions
        from google.cloud.spanner_dbapi.exceptions import OperationalError

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)

        with mock.patch(
            "google.cloud.spanner_dbapi.parse_utils.classify_stmt",
            side_effect=exceptions.InternalServerError("message"),
        ):
            with self.assertRaises(OperationalError):
                cursor.execute(sql="sql")

    def test_executemany_on_closed_cursor(self):
        from google.cloud.spanner_dbapi import InterfaceError
        from google.cloud.spanner_dbapi 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.close()

        with self.assertRaises(InterfaceError):
            cursor.executemany("""SELECT * FROM table1 WHERE "col1" = @a1""", ())

    def test_executemany_DLL(self):
        from google.cloud.spanner_dbapi import connect, ProgrammingError

        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()

        with self.assertRaises(ProgrammingError):
            cursor.executemany("""DROP DATABASE database_name""", ())

    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,)))
        )

    @unittest.skipIf(
        sys.version_info[0] < 3, "Python 2 has an outdated iterator definition"
    )
    def test_fetchone(self):
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor._checksum = ResultsChecksum()
        lst = [1, 2, 3]
        cursor._itr = iter(lst)
        for i in range(len(lst)):
            self.assertEqual(cursor.fetchone(), lst[i])
        self.assertIsNone(cursor.fetchone())

    @unittest.skipIf(
        sys.version_info[0] < 3, "Python 2 has an outdated iterator definition"
    )
    def test_fetchone_w_autocommit(self):
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        connection.autocommit = True
        cursor = self._make_one(connection)
        cursor._checksum = ResultsChecksum()
        lst = [1, 2, 3]
        cursor._itr = iter(lst)
        for i in range(len(lst)):
            self.assertEqual(cursor.fetchone(), lst[i])
        self.assertIsNone(cursor.fetchone())

    def test_fetchmany(self):
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor._checksum = ResultsChecksum()
        lst = [(1,), (2,), (3,)]
        cursor._itr = iter(lst)

        self.assertEqual(cursor.fetchmany(), [lst[0]])

        result = cursor.fetchmany(len(lst))
        self.assertEqual(result, lst[1:])

    def test_fetchmany_w_autocommit(self):
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        connection.autocommit = True
        cursor = self._make_one(connection)
        cursor._checksum = ResultsChecksum()
        lst = [(1,), (2,), (3,)]
        cursor._itr = iter(lst)

        self.assertEqual(cursor.fetchmany(), [lst[0]])

        result = cursor.fetchmany(len(lst))
        self.assertEqual(result, lst[1:])

    def test_fetchall(self):
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor._checksum = ResultsChecksum()
        lst = [(1,), (2,), (3,)]
        cursor._itr = iter(lst)
        self.assertEqual(cursor.fetchall(), lst)

    def test_fetchall_w_autocommit(self):
        from google.cloud.spanner_dbapi.checksum import ResultsChecksum

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        connection.autocommit = True
        cursor = self._make_one(connection)
        cursor._checksum = ResultsChecksum()
        lst = [(1,), (2,), (3,)]
        cursor._itr = iter(lst)
        self.assertEqual(cursor.fetchall(), lst)

    def test_nextset(self):
        from google.cloud.spanner_dbapi import exceptions

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor.close()
        with self.assertRaises(exceptions.InterfaceError):
            cursor.nextset()

    def test_setinputsizes(self):
        from google.cloud.spanner_dbapi import exceptions

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor.close()
        with self.assertRaises(exceptions.InterfaceError):
            cursor.setinputsizes(sizes=None)

    def test_setoutputsize(self):
        from google.cloud.spanner_dbapi import exceptions

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        cursor = self._make_one(connection)
        cursor.close()
        with self.assertRaises(exceptions.InterfaceError):
            cursor.setoutputsize(size=None)

    # def test_handle_insert(self):
    #     pass
    #
    # def test_do_execute_insert_heterogenous(self):
    #     pass
    #
    # def test_do_execute_insert_homogenous(self):
    #     pass

    def test_handle_dql(self):
        from google.cloud.spanner_dbapi import utils
        from google.cloud.spanner_dbapi.cursor import _UNSET_COUNT

        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        connection.database.snapshot.return_value.__enter__.return_value = (
            mock_snapshot
        ) = mock.MagicMock()
        cursor = self._make_one(connection)

        mock_snapshot.execute_sql.return_value = int(0)
        cursor._handle_DQL("sql", params=None)
        self.assertEqual(cursor._row_count, 0)
        self.assertIsNone(cursor._itr)

        mock_snapshot.execute_sql.return_value = "0"
        cursor._handle_DQL("sql", params=None)
        self.assertEqual(cursor._result_set, "0")
        self.assertIsInstance(cursor._itr, utils.PeekIterator)
        self.assertEqual(cursor._row_count, _UNSET_COUNT)

    def test_context(self):
        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        with cursor as c:
            self.assertEqual(c, cursor)

        self.assertTrue(c.is_closed)

    def test_next(self):
        from google.cloud.spanner_dbapi import exceptions

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        with self.assertRaises(exceptions.ProgrammingError):
            cursor.__next__()

        lst = [(1,), (2,), (3,)]
        cursor._itr = iter(lst)
        i = 0
        for c in cursor._itr:
            self.assertEqual(c, lst[i])
            i += 1

    def test_iter(self):
        from google.cloud.spanner_dbapi import exceptions

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)
        with self.assertRaises(exceptions.ProgrammingError):
            _ = iter(cursor)

        iterator = iter([(1,), (2,), (3,)])
        cursor._itr = iterator
        self.assertEqual(iter(cursor), iterator)

    def test_list_tables(self):
        from google.cloud.spanner_dbapi import _helpers

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)

        table_list = ["table1", "table2", "table3"]
        with mock.patch(
            "google.cloud.spanner_dbapi.cursor.Cursor.run_sql_in_snapshot",
            return_value=table_list,
        ) as mock_run_sql:
            cursor.list_tables()
            mock_run_sql.assert_called_once_with(_helpers.SQL_LIST_TABLES)

    def test_run_sql_in_snapshot(self):
        connection = self._make_connection(self.INSTANCE, mock.MagicMock())
        connection.database.snapshot.return_value.__enter__.return_value = (
            mock_snapshot
        ) = mock.MagicMock()
        cursor = self._make_one(connection)

        results = 1, 2, 3
        mock_snapshot.execute_sql.return_value = results
        self.assertEqual(cursor.run_sql_in_snapshot("sql"), list(results))

    def test_get_table_column_schema(self):
        from google.cloud.spanner_dbapi.cursor import ColumnDetails
        from google.cloud.spanner_dbapi import _helpers
        from google.cloud.spanner_v1 import param_types

        connection = self._make_connection(self.INSTANCE, self.DATABASE)
        cursor = self._make_one(connection)

        column_name = "column1"
        is_nullable = "YES"
        spanner_type = "spanner_type"
        rows = [(column_name, is_nullable, spanner_type)]
        expected = {column_name: ColumnDetails(null_ok=True, spanner_type=spanner_type)}
        with mock.patch(
            "google.cloud.spanner_dbapi.cursor.Cursor.run_sql_in_snapshot",
            return_value=rows,
        ) as mock_run_sql:
            table_name = "table1"
            result = cursor.get_table_column_schema(table_name=table_name)
            mock_run_sql.assert_called_once_with(
                sql=_helpers.SQL_GET_TABLE_COLUMN_SCHEMA,
                params={"table_name": table_name},
                param_types={"table_name": param_types.STRING},
            )
            self.assertEqual(result, expected)

    def test_peek_iterator_aborted(self):
        """
        Checking that an Aborted exception is retried in case it happened
        while streaming the first element with a PeekIterator.
        """
        from google.api_core.exceptions import Aborted
        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()
        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),
                ):
                    cursor.execute("SELECT * FROM table_name")

                retry_mock.assert_called_with()

    def test_peek_iterator_aborted_autocommit(self):
        """
        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

        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()
        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_fetchone_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.fetchone()

                retry_mock.assert_called_with()

    def test_fetchone_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)

        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=([row], ResultsChecksum()),
            ) as run_mock:

                cursor.fetchone()

                run_mock.assert_called_with(statement, retried=True)

    def test_fetchone_retry_aborted_statements_checksums_mismatch(self):
        """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"]

        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)

        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.fetchone()

                run_mock.assert_called_with(statement, retried=True)

    def test_fetchall_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.__iter__",
            side_effect=(Aborted("Aborted"), iter([])),
        ):
            with mock.patch(
                "google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
            ) as retry_mock:

                cursor.fetchall()

                retry_mock.assert_called_with()

    def test_fetchall_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)

        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 test_fetchall_retry_aborted_statements_checksums_mismatch(self):
        """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"]

        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)

        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=([row2], ResultsChecksum()),
            ) as run_mock:

                with self.assertRaises(RetryAborted):
                    cursor.fetchall()

                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_fetchmany_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)

        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=([row], ResultsChecksum()),
            ) as run_mock:

                cursor.fetchmany(len(row))

                run_mock.assert_called_with(statement, retried=True)

    def test_fetchmany_retry_aborted_statements_checksums_mismatch(self):
        """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"]

        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)

        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_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",
            (
                "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",
        ]

        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"
            "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)