Repository: googleapis/python-spanner

Issue:
Aborted errors not being retried in DBAPI
From the SQLAlchemy CircleCI:
```
Traceback (most recent call last):
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/testing/suite/test_dialect.py", line 200, in test_turn_autocommit_off_via_default_iso_level
    self._test_conn_autocommits(conn, False)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/testing/suite/test_dialect.py", line 175, in _test_conn_autocommits
    conn.execute(self.tables.some_table.delete())
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/engine/base.py", line 1011, in execute
    return meth(self, multiparams, params)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/sql/elements.py", line 298, in _execute_on_connection
    return connection._execute_clauseelement(self, multiparams, params)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/engine/base.py", line 1124, in _execute_clauseelement
    ret = self._execute_context(
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/engine/base.py", line 1316, in _execute_context
    self._handle_dbapi_exception(
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/engine/base.py", line 1514, in _handle_dbapi_exception
    util.raise_(exc_info[1], with_traceback=exc_info[2])
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/util/compat.py", line 182, in raise_
    raise exception
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/engine/base.py", line 1276, in _execute_context
    self.dialect.do_execute(
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/sqlalchemy/engine/default.py", line 608, in do_execute
    cursor.execute(statement, parameters)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/cloud/spanner_dbapi/cursor.py", line 209, in execute
    self._itr = PeekIterator(self._result_set)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/cloud/spanner_dbapi/utils.py", line 38, in __init__
    head = next(itr_src)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/cloud/spanner_v1/streamed.py", line 149, in __iter__
    self._consume_next()
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/cloud/spanner_v1/streamed.py", line 121, in _consume_next
    response = six.next(self._response_iterator)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/cloud/spanner_v1/snapshot.py", line 58, in _restart_on_unavailable
    iterator = method(request=request)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/cloud/spanner_v1/services/spanner/client.py", line 899, in execute_streaming_sql
    response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/api_core/gapic_v1/method.py", line 145, in __call__
    return wrapped_func(*args, **kwargs)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/api_core/timeout.py", line 102, in func_with_timeout
    return func(*args, **kwargs)
  File "/home/circleci/project/.nox/compliance_test/lib/python3.8/site-packages/google/api_core/grpc_helpers.py", line 168, in error_remapped_callable
    six.raise_from(exceptions.from_grpc_error(exc), exc)
  File "<string>", line 3, in raise_from
google.api_core.exceptions.Aborted: 409 Transaction aborted. Database schema probably changed during transaction, retry may succeed.
```

DBAPI only retries `Aborted` errors in the fetch functions. However, the `Aborted` error is being raised in `execute()` by the `PeekIterator`. I think the `PeekIterator` implementation needs to be updated to prevent iterating before the fetch functions.

Discussion:

Patch:
diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py
index a28879f..3569bab 100644
--- a/google/cloud/spanner_dbapi/cursor.py
+++ b/google/cloud/spanner_dbapi/cursor.py
@@ -206,7 +206,12 @@ class Cursor(object):
                 (self._result_set, self._checksum,) = self.connection.run_statement(
                     statement
                 )
-                self._itr = PeekIterator(self._result_set)
+                while True:
+                    try:
+                        self._itr = PeekIterator(self._result_set)
+                        break
+                    except Aborted:
+                        self.connection.retry_transaction()
                 return
 
             if classification == parse_utils.STMT_NON_UPDATING:
@@ -352,7 +357,12 @@ class Cursor(object):
                 self._result_set = res
                 # Read the first element so that the StreamedResultSet can
                 # return the metadata after a DQL statement. See issue #155.
-                self._itr = PeekIterator(self._result_set)
+                while True:
+                    try:
+                        self._itr = PeekIterator(self._result_set)
+                        break
+                    except Aborted:
+                        self.connection.retry_transaction()
                 # Unfortunately, Spanner doesn't seem to send back
                 # information about the number of rows available.
                 self._row_count = _UNSET_COUNT
diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py
index a9ae36d..1385809 100644
--- a/google/cloud/spanner_v1/_helpers.py
+++ b/google/cloud/spanner_v1/_helpers.py
@@ -30,6 +30,16 @@ from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 
 
+# Validation error messages
+NUMERIC_MAX_SCALE_ERR_MSG = (
+    "Max scale for a numeric is 9. The requested numeric has scale {}"
+)
+NUMERIC_MAX_PRECISION_ERR_MSG = (
+    "Max precision for the whole component of a numeric is 29. The requested "
+    + "numeric has a whole component with precision {}"
+)
+
+
 def _try_to_coerce_bytes(bytestring):
     """Try to coerce a byte string into the right thing based on Python
     version and whether or not it is base64 encoded.
@@ -87,6 +97,28 @@ def _merge_query_options(base, merge):
     return combined
 
 
+def _assert_numeric_precision_and_scale(value):
+    """
+    Asserts that input numeric field is within Spanner supported range.
+
+    Spanner supports fixed 38 digits of precision and 9 digits of scale.
+    This number can be optionally prefixed with a plus or minus sign.
+    Read more here: https://cloud.google.com/spanner/docs/data-types#numeric_type
+
+    :type value: decimal.Decimal
+    :param value: The value to check for Cloud Spanner compatibility.
+
+    :raises NotSupportedError: If value is not within supported precision or scale of Spanner.
+    """
+    scale = value.as_tuple().exponent
+    precision = len(value.as_tuple().digits)
+
+    if scale < -9:
+        raise ValueError(NUMERIC_MAX_SCALE_ERR_MSG.format(abs(scale)))
+    if precision + scale > 29:
+        raise ValueError(NUMERIC_MAX_PRECISION_ERR_MSG.format(precision + scale))
+
+
 # pylint: disable=too-many-return-statements,too-many-branches
 def _make_value_pb(value):
     """Helper for :func:`_make_list_value_pbs`.
@@ -129,6 +161,7 @@ def _make_value_pb(value):
     if isinstance(value, ListValue):
         return Value(list_value=value)
     if isinstance(value, decimal.Decimal):
+        _assert_numeric_precision_and_scale(value)
         return Value(string_value=str(value))
     raise ValueError("Unknown type: %s" % (value,))
 


Test patch:
diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py
index 4d5db01..57a3375 100644
--- a/tests/unit/spanner_dbapi/test_cursor.py
+++ b/tests/unit/spanner_dbapi/test_cursor.py
@@ -549,6 +549,74 @@ class TestCursor(unittest.TestCase):
             )
             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
diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py
index 73277a7..11239d7 100644
--- a/tests/unit/spanner_dbapi/test_parse_utils.py
+++ b/tests/unit/spanner_dbapi/test_parse_utils.py
@@ -254,8 +254,6 @@ class TestParseUtils(unittest.TestCase):
 
     @unittest.skipIf(skip_condition, skip_message)
     def test_sql_pyformat_args_to_spanner(self):
-        import decimal
-
         from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner
 
         cases = [
@@ -300,16 +298,6 @@ class TestParseUtils(unittest.TestCase):
                 ("SELECT * from t WHERE id=10", {"f1": "app", "f2": "name"}),
                 ("SELECT * from t WHERE id=10", {"f1": "app", "f2": "name"}),
             ),
-            (
-                (
-                    "SELECT (an.p + %s) AS np FROM an WHERE (an.p + %s) = %s",
-                    (1, 1.0, decimal.Decimal("31")),
-                ),
-                (
-                    "SELECT (an.p + @a0) AS np FROM an WHERE (an.p + @a1) = @a2",
-                    {"a0": 1, "a1": 1.0, "a2": decimal.Decimal("31")},
-                ),
-            ),
         ]
         for ((sql_in, params), sql_want) in cases:
             with self.subTest(sql=sql_in):
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
index fecf258..305a6ce 100644
--- a/tests/unit/test__helpers.py
+++ b/tests/unit/test__helpers.py
@@ -233,6 +233,65 @@ class Test_make_value_pb(unittest.TestCase):
         with self.assertRaises(ValueError):
             self._callFUT(object())
 
+    def test_w_numeric_precision_and_scale_valid(self):
+        import decimal
+        from google.protobuf.struct_pb2 import Value
+
+        cases = [
+            decimal.Decimal("42"),
+            decimal.Decimal("9.9999999999999999999999999999999999999E+28"),
+            decimal.Decimal("-9.9999999999999999999999999999999999999E+28"),
+            decimal.Decimal("99999999999999999999999999999.999999999"),
+            decimal.Decimal("1E+28"),
+            decimal.Decimal("1E-9"),
+        ]
+        for value in cases:
+            with self.subTest(value=value):
+                value_pb = self._callFUT(value)
+                self.assertIsInstance(value_pb, Value)
+                self.assertEqual(value_pb.string_value, str(value))
+
+    def test_w_numeric_precision_and_scale_invalid(self):
+        import decimal
+        from google.cloud.spanner_v1._helpers import (
+            NUMERIC_MAX_SCALE_ERR_MSG,
+            NUMERIC_MAX_PRECISION_ERR_MSG,
+        )
+
+        max_precision_error_msg = NUMERIC_MAX_PRECISION_ERR_MSG.format("30")
+        max_scale_error_msg = NUMERIC_MAX_SCALE_ERR_MSG.format("10")
+
+        cases = [
+            (
+                decimal.Decimal("9.9999999999999999999999999999999999999E+29"),
+                max_precision_error_msg,
+            ),
+            (
+                decimal.Decimal("-9.9999999999999999999999999999999999999E+29"),
+                max_precision_error_msg,
+            ),
+            (
+                decimal.Decimal("999999999999999999999999999999.99999999"),
+                max_precision_error_msg,
+            ),
+            (
+                decimal.Decimal("-999999999999999999999999999999.99999999"),
+                max_precision_error_msg,
+            ),
+            (
+                decimal.Decimal("999999999999999999999999999999"),
+                max_precision_error_msg,
+            ),
+            (decimal.Decimal("1E+29"), max_precision_error_msg),
+            (decimal.Decimal("1E-10"), max_scale_error_msg),
+        ]
+
+        for value, err_msg in cases:
+            with self.subTest(value=value, err_msg=err_msg):
+                self.assertRaisesRegex(
+                    ValueError, err_msg, lambda: self._callFUT(value),
+                )
+
 
 class Test_make_list_value_pb(unittest.TestCase):
     def _callFUT(self, *args, **kw):