google-cloud-spanner 3.55.0__py3-none-any.whl

google/cloud/spanner_dbapi/parser.py

@@ -0,0 +1,258 @@
+# 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.
+
+"""
+Grammar for parsing VALUES:
+    VALUES      := `VALUES(` + ARGS + `)`
+    ARGS        := [EXPR,]*EXPR
+    EXPR        := TERMINAL / FUNC
+    TERMINAL    := `%s`
+    FUNC        := alphanum + `(` + ARGS + `)`
+    alphanum    := (a-zA-Z_)[0-9a-ZA-Z_]*
+
+thus given:
+    statement: 'VALUES (%s, %s), (%s, LOWER(UPPER(%s)))   , (%s)'
+    It'll parse:
+        VALUES
+            |- ARGS
+                |- (TERMINAL, TERMINAL)
+                |- (TERMINAL, FUNC
+                                |- FUNC
+                                    |- (TERMINAL)
+                |- (TERMINAL)
+"""
+
+from .exceptions import ProgrammingError
+
+ARGS = "ARGS"
+FUNC = "FUNC"
+VALUES = "VALUES"
+
+
+class func(object):
+    def __init__(self, func_name, args):
+        self.name = func_name
+        self.args = args
+
+    def __str__(self):
+        return "%s%s" % (self.name, self.args)
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+        if self.name != other.name:
+            return False
+        if not isinstance(other.args, type(self.args)):
+            return False
+        if len(self.args) != len(other.args):
+            return False
+        return self.args == other.args
+
+    def __len__(self):
+        return len(self.args)
+
+
+class terminal(str):
+    """Represent the unit symbol that can be part of a SQL values clause."""
+
+    pass
+
+
+class a_args(object):
+    """Expression arguments.
+
+    :type argv: list
+    :param argv: A List of expression arguments.
+    """
+
+    def __init__(self, argv):
+        self.argv = argv
+
+    def __str__(self):
+        return "(" + ", ".join([str(arg) for arg in self.argv]) + ")"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def has_expr(self):
+        return any([token for token in self.argv if not isinstance(token, terminal)])
+
+    def __len__(self):
+        return len(self.argv)
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        if len(self) != len(other):
+            return False
+
+        for i, item in enumerate(self):
+            if item != other[i]:
+                return False
+
+        return True
+
+    def __getitem__(self, index):
+        return self.argv[index]
+
+    def homogenous(self):
+        """Check arguments of the expression to be homogeneous.
+
+        :rtype: bool
+        :return: True if all the arguments of the expression are in pyformat
+                 and each has the same length, False otherwise.
+        """
+        if not self._is_equal_length():
+            return False
+
+        for arg in self.argv:
+            if isinstance(arg, terminal):
+                continue
+            elif isinstance(arg, a_args):
+                if not arg.homogenous():
+                    return False
+            else:
+                return False
+        return True
+
+    def _is_equal_length(self):
+        """Return False if all the arguments have the same length.
+
+        :rtype: bool
+        :return: False if the sequences of the arguments have the same length.
+        """
+        if len(self) == 0:
+            return True
+
+        arg0_len = len(self.argv[0])
+        for arg in self.argv[1:]:
+            if len(arg) != arg0_len:
+                return False
+
+        return True
+
+
+class values(a_args):
+    """A wrapper for values.
+
+    :rtype: str
+    :returns: A string of the values expression in a tree view.
+    """
+
+    def __str__(self):
+        return "VALUES%s" % super().__str__()
+
+
+pyfmt_str = terminal("%s")
+
+
+def expect(word, token):
+    """Parse the given expression recursively.
+
+    :type word: str
+    :param word: A string expression.
+
+    :type token: str
+    :param token: An expression token.
+
+    :rtype: `Tuple(str, Any)`
+    :returns: A tuple containing the rest of the expression string and the
+              parse tree for the part of the expression that has already been
+              parsed.
+
+    :raises :class:`ProgrammingError`: If there is a parsing error.
+    """
+    word = word.strip()
+    if token == VALUES:
+        if not word.startswith("VALUES"):
+            raise ProgrammingError("VALUES: `%s` does not start with VALUES" % word)
+
+        word = word[len("VALUES") :].lstrip()
+
+        all_args = []
+        while word:
+            word = word.strip()
+
+            word, arg = expect(word, ARGS)
+            all_args.append(arg)
+            word = word.strip()
+
+            if word and not word.startswith(","):
+                raise ProgrammingError(
+                    "VALUES: expected `,` got %s in %s" % (word[0], word)
+                )
+            word = word[1:]
+        return "", values(all_args)
+
+    elif token == FUNC:
+        begins_with_letter = word and (word[0].isalpha() or word[0] == "_")
+        if not begins_with_letter:
+            raise ProgrammingError(
+                "FUNC: `%s` does not begin with `a-zA-z` nor a `_`" % word
+            )
+
+        rest = word[1:]
+        end = 0
+        for ch in rest:
+            if ch.isalnum() or ch == "_":
+                end += 1
+            else:
+                break
+
+        func_name, rest = word[: end + 1], word[end + 1 :].strip()
+
+        word, args = expect(rest, ARGS)
+        return word, func(func_name, args)
+
+    elif token == ARGS:
+        # The form should be:
+        #   (%s)
+        #   (%s, %s...)
+        #   (FUNC, %s...)
+        #   (%s, %s...)
+        if not (word and word.startswith("(")):
+            raise ProgrammingError("ARGS: supposed to begin with `(` in `%s`" % word)
+
+        word = word[1:]
+
+        terms = []
+        while True:
+            word = word.strip()
+            if not word or word.startswith(")"):
+                break
+
+            if word == "%s":
+                terms.append(pyfmt_str)
+                word = ""
+            elif not word.startswith("%s"):
+                word, parsed = expect(word, FUNC)
+                terms.append(parsed)
+            else:
+                terms.append(pyfmt_str)
+                word = word[2:].strip()
+
+            if word.startswith(","):
+                word = word[1:]
+
+        if not (word and word.startswith(")")):
+            raise ProgrammingError("ARGS: supposed to end with `)` in `%s`" % word)
+
+        word = word[1:]
+        return word, a_args(terms)
+
+    raise ProgrammingError("Unknown token `%s`" % token)

google/cloud/spanner_dbapi/partition_helper.py

@@ -0,0 +1,41 @@
+# Copyright 2023 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.
+
+from dataclasses import dataclass
+from typing import Any
+
+import gzip
+import pickle
+import base64
+
+from google.cloud.spanner_v1 import BatchTransactionId
+
+
+def decode_from_string(encoded_partition_id):
+    gzip_bytes = base64.b64decode(bytes(encoded_partition_id, "utf-8"))
+    partition_id_bytes = gzip.decompress(gzip_bytes)
+    return pickle.loads(partition_id_bytes)
+
+
+def encode_to_string(batch_transaction_id, partition_result):
+    partition_id = PartitionId(batch_transaction_id, partition_result)
+    partition_id_bytes = pickle.dumps(partition_id)
+    gzip_bytes = gzip.compress(partition_id_bytes)
+    return str(base64.b64encode(gzip_bytes), "utf-8")
+
+
+@dataclass
+class PartitionId:
+    batch_transaction_id: BatchTransactionId
+    partition_result: Any

google/cloud/spanner_dbapi/transaction_helper.py

@@ -0,0 +1,294 @@
+# Copyright 2023 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.
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, List, Any, Dict
+from google.api_core.exceptions import Aborted
+
+import time
+
+from google.cloud.spanner_dbapi.batch_dml_executor import BatchMode
+from google.cloud.spanner_dbapi.exceptions import RetryAborted
+from google.cloud.spanner_v1._helpers import _get_retry_delay
+
+if TYPE_CHECKING:
+    from google.cloud.spanner_dbapi import Connection, Cursor
+from google.cloud.spanner_dbapi.checksum import ResultsChecksum, _compare_checksums
+
+MAX_INTERNAL_RETRIES = 50
+RETRY_ABORTED_ERROR = "The transaction was aborted and could not be retried due to a concurrent modification."
+
+
+class TransactionRetryHelper:
+    def __init__(self, connection: "Connection"):
+        """Helper class used in retrying the transaction when aborted This will
+        maintain all the statements executed on original transaction and replay
+        them again in the retried transaction.
+
+        :type connection: :class:`~google.cloud.spanner_dbapi.connection.Connection`
+        :param connection: A DB-API connection to Google Cloud Spanner.
+        """
+
+        self._connection = connection
+        # list of all statements in the same order as executed in original
+        # transaction along with their results
+        self._statement_result_details_list: List[StatementDetails] = []
+        # Map of last StatementDetails that was added to a particular cursor
+        self._last_statement_details_per_cursor: Dict[Cursor, StatementDetails] = {}
+        # 1-1 map from original cursor object on which transaction ran to the
+        # new cursor object used in the retry
+        self._cursor_map: Dict[Cursor, Cursor] = {}
+
+    def _set_connection_for_retry(self):
+        self._connection._spanner_transaction_started = False
+        self._connection._transaction_begin_marked = False
+        self._connection._batch_mode = BatchMode.NONE
+
+    def reset(self):
+        """
+        Resets the state of the class when the ongoing transaction is committed
+        or aborted
+        """
+        self._statement_result_details_list = []
+        self._last_statement_details_per_cursor = {}
+        self._cursor_map = {}
+
+    def add_fetch_statement_for_retry(
+        self, cursor, result_rows, exception, is_fetch_all
+    ):
+        """
+        StatementDetails to be added to _statement_result_details_list whenever fetchone, fetchmany or
+        fetchall method is called on the cursor.
+        If fetchone is consecutively called n times then it is stored as fetchmany with size as n.
+        Same for fetchmany, so consecutive fetchone and fetchmany statements are stored as one
+        fetchmany statement in _statement_result_details_list with size param appropriately set
+
+        :param cursor: original Cursor object on which statement executed in the transaction
+        :param result_rows: All the rows from the resultSet from fetch statement execution
+        :param exception: Not none in case non-aborted exception is thrown on the original
+        statement execution
+        :param is_fetch_all: True in case of fetchall statement execution
+        """
+        if not self._connection._client_transaction_started:
+            return
+
+        last_statement_result_details = self._last_statement_details_per_cursor.get(
+            cursor
+        )
+        if (
+            last_statement_result_details is not None
+            and last_statement_result_details.statement_type
+            == CursorStatementType.FETCH_MANY
+        ):
+            if exception is not None:
+                last_statement_result_details.result_type = ResultType.EXCEPTION
+                last_statement_result_details.result_details = exception
+            else:
+                for row in result_rows:
+                    last_statement_result_details.result_details.consume_result(row)
+                last_statement_result_details.size += len(result_rows)
+        else:
+            result_details = _get_statement_result_checksum(result_rows)
+            if is_fetch_all:
+                statement_type = CursorStatementType.FETCH_ALL
+                size = None
+            else:
+                statement_type = CursorStatementType.FETCH_MANY
+                size = len(result_rows)
+
+            last_statement_result_details = FetchStatement(
+                cursor=cursor,
+                statement_type=statement_type,
+                result_type=ResultType.CHECKSUM,
+                result_details=result_details,
+                size=size,
+            )
+            self._last_statement_details_per_cursor[
+                cursor
+            ] = last_statement_result_details
+            self._statement_result_details_list.append(last_statement_result_details)
+
+    def add_execute_statement_for_retry(
+        self, cursor, sql, args, exception, is_execute_many
+    ):
+        """
+        StatementDetails to be added to _statement_result_details_list whenever execute or
+        executemany method is called on the cursor.
+
+        :param cursor: original Cursor object on which statement executed in the transaction
+        :param sql: Input param of the execute/executemany method
+        :param args: Input param of the execute/executemany method
+        :param exception: Not none in case non-aborted exception is thrown on the original
+        statement execution
+        :param is_execute_many: True in case of executemany statement execution
+        """
+        if not self._connection._client_transaction_started:
+            return
+        statement_type = CursorStatementType.EXECUTE
+        if is_execute_many:
+            statement_type = CursorStatementType.EXECUTE_MANY
+
+        result_type = ResultType.NONE
+        result_details = None
+        if exception is not None:
+            result_type = ResultType.EXCEPTION
+            result_details = exception
+        elif cursor._batch_dml_rows_count is not None:
+            result_type = ResultType.BATCH_DML_ROWS_COUNT
+            result_details = cursor._batch_dml_rows_count
+        elif cursor._row_count is not None:
+            result_type = ResultType.ROW_COUNT
+            result_details = cursor.rowcount
+
+        last_statement_result_details = ExecuteStatement(
+            cursor=cursor,
+            statement_type=statement_type,
+            sql=sql,
+            args=args,
+            result_type=result_type,
+            result_details=result_details,
+        )
+        self._last_statement_details_per_cursor[cursor] = last_statement_result_details
+        self._statement_result_details_list.append(last_statement_result_details)
+
+    def retry_transaction(self, default_retry_delay=None):
+        """Retry the aborted transaction.
+
+        All the statements executed in the original transaction
+        will be re-executed in new one. Results checksums of the
+        original statements and the retried ones will be compared.
+
+        :raises: :class:`google.cloud.spanner_dbapi.exceptions.RetryAborted`
+            If results checksum of the retried statement is
+            not equal to the checksum of the original one.
+        """
+        attempt = 0
+        while True:
+            attempt += 1
+            if attempt > MAX_INTERNAL_RETRIES:
+                raise
+            self._set_connection_for_retry()
+            try:
+                for statement_result_details in self._statement_result_details_list:
+                    if statement_result_details.cursor in self._cursor_map:
+                        cursor = self._cursor_map.get(statement_result_details.cursor)
+                    else:
+                        cursor = self._connection.cursor()
+                        cursor._in_retry_mode = True
+                        self._cursor_map[statement_result_details.cursor] = cursor
+                    try:
+                        _handle_statement(statement_result_details, cursor)
+                    except Aborted:
+                        raise
+                    except RetryAborted:
+                        raise
+                    except Exception as ex:
+                        if (
+                            type(statement_result_details.result_details)
+                            is not type(ex)
+                            or ex.args != statement_result_details.result_details.args
+                        ):
+                            raise RetryAborted(RETRY_ABORTED_ERROR, ex)
+                return
+            except Aborted as ex:
+                delay = _get_retry_delay(
+                    ex.errors[0], attempt, default_retry_delay=default_retry_delay
+                )
+                if delay:
+                    time.sleep(delay)
+
+
+def _handle_statement(statement_result_details, cursor):
+    statement_type = statement_result_details.statement_type
+    if _is_execute_type_statement(statement_type):
+        if statement_type == CursorStatementType.EXECUTE:
+            cursor.execute(statement_result_details.sql, statement_result_details.args)
+            if (
+                statement_result_details.result_type == ResultType.ROW_COUNT
+                and statement_result_details.result_details != cursor.rowcount
+            ):
+                raise RetryAborted(RETRY_ABORTED_ERROR)
+        else:
+            cursor.executemany(
+                statement_result_details.sql, statement_result_details.args
+            )
+        if (
+            statement_result_details.result_type == ResultType.BATCH_DML_ROWS_COUNT
+            and statement_result_details.result_details != cursor._batch_dml_rows_count
+        ):
+            raise RetryAborted(RETRY_ABORTED_ERROR)
+    else:
+        if statement_type == CursorStatementType.FETCH_ALL:
+            res = cursor.fetchall()
+        else:
+            res = cursor.fetchmany(statement_result_details.size)
+        checksum = _get_statement_result_checksum(res)
+        _compare_checksums(checksum, statement_result_details.result_details)
+    if statement_result_details.result_type == ResultType.EXCEPTION:
+        raise RetryAborted(RETRY_ABORTED_ERROR)
+
+
+def _is_execute_type_statement(statement_type):
+    return statement_type in (
+        CursorStatementType.EXECUTE,
+        CursorStatementType.EXECUTE_MANY,
+    )
+
+
+def _get_statement_result_checksum(res_iter):
+    retried_checksum = ResultsChecksum()
+    for res in res_iter:
+        retried_checksum.consume_result(res)
+    return retried_checksum
+
+
+class CursorStatementType(Enum):
+    EXECUTE = 1
+    EXECUTE_MANY = 2
+    FETCH_ONE = 3
+    FETCH_ALL = 4
+    FETCH_MANY = 5
+
+
+class ResultType(Enum):
+    # checksum of ResultSet in case of fetch call on query statement
+    CHECKSUM = 1
+    # None in case of execute call on query statement
+    NONE = 2
+    # Exception details in case of any statement execution throws exception
+    EXCEPTION = 3
+    # Total rows updated in case of execute call on DML statement
+    ROW_COUNT = 4
+    # Total rows updated in case of Batch DML statement execution
+    BATCH_DML_ROWS_COUNT = 5
+
+
+@dataclass
+class StatementDetails:
+    statement_type: CursorStatementType
+    # The cursor object on which this statement was executed
+    cursor: "Cursor"
+    result_type: ResultType
+    result_details: Any
+
+
+@dataclass
+class ExecuteStatement(StatementDetails):
+    sql: str
+    args: Any = None
+
+
+@dataclass
+class FetchStatement(StatementDetails):
+    size: int = None

google/cloud/spanner_dbapi/types.py

@@ -0,0 +1,106 @@
+# 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.
+
+"""Implementation of the type objects and constructors according to the
+   PEP-0249 specification.
+
+   See
+   https://www.python.org/dev/peps/pep-0249/#type-objects-and-constructors
+"""
+
+import datetime
+import time
+from base64 import b64encode
+
+
+def _date_from_ticks(ticks):
+    """Based on PEP-249 Implementation Hints for Module Authors:
+
+    https://www.python.org/dev/peps/pep-0249/#implementation-hints-for-module-authors
+    """
+    return Date(*time.localtime(ticks)[:3])
+
+
+def _time_from_ticks(ticks):
+    """Based on PEP-249 Implementation Hints for Module Authors:
+
+    https://www.python.org/dev/peps/pep-0249/#implementation-hints-for-module-authors
+    """
+    return Time(*time.localtime(ticks)[3:6])
+
+
+def _timestamp_from_ticks(ticks):
+    """Based on PEP-249 Implementation Hints for Module Authors:
+
+    https://www.python.org/dev/peps/pep-0249/#implementation-hints-for-module-authors
+    """
+    return Timestamp(*time.localtime(ticks)[:6])
+
+
+class _DBAPITypeObject(object):
+    """Implementation of a helper class used for type comparison among similar
+    but possibly different types.
+
+    See
+    https://www.python.org/dev/peps/pep-0249/#implementation-hints-for-module-authors
+    """
+
+    def __init__(self, *values):
+        self.values = values
+
+    def __eq__(self, other):
+        return other in self.values
+
+
+Date = datetime.date
+Time = datetime.time
+Timestamp = datetime.datetime
+DateFromTicks = _date_from_ticks
+TimeFromTicks = _time_from_ticks
+TimestampFromTicks = _timestamp_from_ticks
+Binary = b64encode
+
+STRING = "STRING"
+BINARY = _DBAPITypeObject("TYPE_CODE_UNSPECIFIED", "BYTES", "ARRAY", "STRUCT")
+NUMBER = _DBAPITypeObject("BOOL", "INT64", "FLOAT64", "FLOAT32", "NUMERIC")
+DATETIME = _DBAPITypeObject("TIMESTAMP", "DATE")
+ROWID = "STRING"
+
+
+class TimestampStr(str):
+    """[inherited from the alpha release]
+
+    TODO: Decide whether this class is necessary
+
+    TimestampStr exists so that we can purposefully format types as timestamps
+    compatible with Cloud Spanner's TIMESTAMP type, but right before making
+    queries, it'll help differentiate between normal strings and the case of
+    types that should be TIMESTAMP.
+    """
+
+    pass
+
+
+class DateStr(str):
+    """[inherited from the alpha release]
+
+    TODO: Decide whether this class is necessary
+
+    DateStr is a sentinel type to help format Django dates as
+    compatible with Cloud Spanner's DATE type, but right before making
+    queries, it'll help differentiate between normal strings and the case of
+    types that should be DATE.
+    """
+
+    pass