feat: Implementation of client side statements that return (#1046)

* Implementation of client side statements that return * Small fix * Incorporated comments * Added tests for exception in commit and rollback * Fix in tests * Skipping few tests from running in emulator * Few fixes * Refactoring * Incorporated comments * Incorporating comments

Author: ankiaga

google/cloud/spanner_dbapi/client_side_statement_executor.py

@@ -15,10 +15,27 @@
 
 if TYPE_CHECKING:
  from google.cloud.spanner_dbapi import Connection
+ from google.cloud.spanner_dbapi import ProgrammingError
+
 from google.cloud.spanner_dbapi.parsed_statement import (
  ParsedStatement,
  ClientSideStatementType,
 )
+from google.cloud.spanner_v1 import (
+ Type,
+ StructType,
+ TypeCode,
+ ResultSetMetadata,
+ PartialResultSet,
+)
+
+from google.cloud.spanner_v1._helpers import _make_value_pb
+from google.cloud.spanner_v1.streamed import StreamedResultSet
+
+CONNECTION_CLOSED_ERROR = "This connection is closed"
+TRANSACTION_NOT_STARTED_WARNING = (
+ "This method is non-operational as a transaction has not been started."
+)
 
 
 def execute(connection: "Connection", parsed_statement: ParsedStatement):
@@ -32,9 +49,46 @@ def execute(connection: "Connection", parsed_statement: ParsedStatement):
  :type parsed_statement: ParsedStatement
  :param parsed_statement: parsed_statement based on the sql query
  """
- if parsed_statement.client_side_statement_type == ClientSideStatementType.COMMIT:
- return connection.commit()
- if parsed_statement.client_side_statement_type == ClientSideStatementType.BEGIN:
- return connection.begin()
- if parsed_statement.client_side_statement_type == ClientSideStatementType.ROLLBACK:
- return connection.rollback()
+ if connection.is_closed:
+ raise ProgrammingError(CONNECTION_CLOSED_ERROR)
+ statement_type = parsed_statement.client_side_statement_type
+ if statement_type == ClientSideStatementType.COMMIT:
+ connection.commit()
+ return None
+ if statement_type == ClientSideStatementType.BEGIN:
+ connection.begin()
+ return None
+ if statement_type == ClientSideStatementType.ROLLBACK:
+ connection.rollback()
+ return None
+ if statement_type == ClientSideStatementType.SHOW_COMMIT_TIMESTAMP:
+ if connection._transaction is None:
+ committed_timestamp = None
+ else:
+ committed_timestamp = connection._transaction.committed
+ return _get_streamed_result_set(
+ ClientSideStatementType.SHOW_COMMIT_TIMESTAMP.name,
+ TypeCode.TIMESTAMP,
+ committed_timestamp,
+ )
+ if statement_type == ClientSideStatementType.SHOW_READ_TIMESTAMP:
+ if connection._snapshot is None:
+ read_timestamp = None
+ else:
+ read_timestamp = connection._snapshot._transaction_read_timestamp
+ return _get_streamed_result_set(
+ ClientSideStatementType.SHOW_READ_TIMESTAMP.name,
+ TypeCode.TIMESTAMP,
+ read_timestamp,
+ )
+
+
+def _get_streamed_result_set(column_name, type_code, column_value):
+ struct_type_pb = StructType(
+ fields=[StructType.Field(name=column_name, type_=Type(code=type_code))]
+ )
+
+ result_set = PartialResultSet(metadata=ResultSetMetadata(row_type=struct_type_pb))
+ if column_value is not None:
+ result_set.values.extend([_make_value_pb(column_value)])
+ return StreamedResultSet(iter([result_set]))

google/cloud/spanner_dbapi/client_side_statement_parser.py

@@ -23,6 +23,12 @@
 RE_BEGIN = re.compile(r"^\s*(BEGIN|START)(TRANSACTION)?", re.IGNORECASE)
 RE_COMMIT = re.compile(r"^\s*(COMMIT)(TRANSACTION)?", re.IGNORECASE)
 RE_ROLLBACK = re.compile(r"^\s*(ROLLBACK)(TRANSACTION)?", re.IGNORECASE)
+RE_SHOW_COMMIT_TIMESTAMP = re.compile(
+ r"^\s*(SHOW)\s+(VARIABLE)\s+(COMMIT_TIMESTAMP)", re.IGNORECASE
+)
+RE_SHOW_READ_TIMESTAMP = re.compile(
+ r"^\s*(SHOW)\s+(VARIABLE)\s+(READ_TIMESTAMP)", re.IGNORECASE
+)
 
 
 def parse_stmt(query):
@@ -37,16 +43,19 @@ def parse_stmt(query):
  :rtype: ParsedStatement
  :returns: ParsedStatement object.
  """
+ client_side_statement_type = None
  if RE_COMMIT.match(query):
- return ParsedStatement(
- StatementType.CLIENT_SIDE, query, ClientSideStatementType.COMMIT
- )
+ client_side_statement_type = ClientSideStatementType.COMMIT
  if RE_BEGIN.match(query):
- return ParsedStatement(
- StatementType.CLIENT_SIDE, query, ClientSideStatementType.BEGIN
- )
+ client_side_statement_type = ClientSideStatementType.BEGIN
  if RE_ROLLBACK.match(query):
+ client_side_statement_type = ClientSideStatementType.ROLLBACK
+ if RE_SHOW_COMMIT_TIMESTAMP.match(query):
+ client_side_statement_type = ClientSideStatementType.SHOW_COMMIT_TIMESTAMP
+ if RE_SHOW_READ_TIMESTAMP.match(query):
+ client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP
+ if client_side_statement_type is not None:
  return ParsedStatement(
- StatementType.CLIENT_SIDE, query, ClientSideStatementType.ROLLBACK
+ StatementType.CLIENT_SIDE, query, client_side_statement_type
  )
  return None

google/cloud/spanner_dbapi/connection.py

@@ -23,6 +23,7 @@
 from google.cloud.spanner_v1 import RequestOptions
 from google.cloud.spanner_v1.session import _get_retry_delay
 from google.cloud.spanner_v1.snapshot import Snapshot
+from deprecated import deprecated
 
 from google.cloud.spanner_dbapi.checksum import _compare_checksums
 from google.cloud.spanner_dbapi.checksum import ResultsChecksum
@@ -35,7 +36,7 @@
 
 
 CLIENT_TRANSACTION_NOT_STARTED_WARNING = (
- "This method is non-operational as transaction has not started"
+ "This method is non-operational as a transaction has not been started."
 )
 MAX_INTERNAL_RETRIES = 50
 
@@ -107,6 +108,9 @@ def __init__(self, instance, database=None, read_only=False):
  self._staleness = None
  self.request_priority = None
  self._transaction_begin_marked = False
+ # whether transaction started at Spanner. This means that we had
+ # made atleast one call to Spanner.
+ self._spanner_transaction_started = False
 
  @property
  def autocommit(self):
@@ -140,26 +144,15 @@ def database(self):
  return self._database
 
  @property
- def _spanner_transaction_started(self):
- """Flag: whether transaction started at Spanner. This means that we had
- made atleast one call to Spanner. Property client_transaction_started
- would always be true if this is true as transaction has to start first
- at clientside than at Spanner
-
- Returns:
- bool: True if Spanner transaction started, False otherwise.
- """
+ @deprecated(
+ reason="This method is deprecated. Use _spanner_transaction_started field"
+ )
+ def inside_transaction(self):
  return (
  self._transaction
  and not self._transaction.committed
  and not self._transaction.rolled_back
- ) or (self._snapshot is not None)
-
- @property
- def inside_transaction(self):
- """Deprecated property which won't be supported in future versions.
- Please use spanner_transaction_started property instead."""
- return self._spanner_transaction_started
+ )
 
  @property
  def _client_transaction_started(self):
@@ -277,7 +270,8 @@ def _release_session(self):
  """
  if self.database is None:
  raise ValueError("Database needs to be passed for this operation")
- self.database._pool.put(self._session)
+ if self._session is not None:
+ self.database._pool.put(self._session)
  self._session = None
 
  def retry_transaction(self):
@@ -293,7 +287,7 @@ def retry_transaction(self):
  """
  attempt = 0
  while True:
- self._transaction = None
+ self._spanner_transaction_started = False
  attempt += 1
  if attempt > MAX_INTERNAL_RETRIES:
  raise
@@ -319,7 +313,6 @@ def _rerun_previous_statements(self):
  status, res = transaction.batch_update(statements)
 
  if status.code == ABORTED:
- self.connection._transaction = None
  raise Aborted(status.details)
 
  retried_checksum = ResultsChecksum()
@@ -363,6 +356,8 @@ def transaction_checkout(self):
  if not self.read_only and self._client_transaction_started:
  if not self._spanner_transaction_started:
  self._transaction = self._session_checkout().transaction()
+ self._snapshot = None
+ self._spanner_transaction_started = True
  self._transaction.begin()
 
  return self._transaction
@@ -377,11 +372,13 @@ def snapshot_checkout(self):
  :returns: A Cloud Spanner snapshot object, ready to use.
  """
  if self.read_only and self._client_transaction_started:
- if not self._snapshot:
+ if not self._spanner_transaction_started:
  self._snapshot = Snapshot(
  self._session_checkout(), multi_use=True, **self.staleness
  )
+ self._transaction = None
  self._snapshot.begin()
+ self._spanner_transaction_started = True
 
  return self._snapshot
 
@@ -391,7 +388,7 @@ def close(self):
  The connection will be unusable from this point forward. If the
  connection has an active transaction, it will be rolled back.
  """
- if self._spanner_transaction_started and not self.read_only:
+ if self._spanner_transaction_started and not self._read_only:
  self._transaction.rollback()
 
  if self._own_pool and self.database:
@@ -405,13 +402,15 @@ def begin(self):
  Marks the transaction as started.
 
  :raises: :class:`InterfaceError`: if this connection is closed.
- :raises: :class:`OperationalError`: if there is an existing transaction that has begin or is running
+ :raises: :class:`OperationalError`: if there is an existing transaction
+ that has been started
  """
  if self._transaction_begin_marked:
  raise OperationalError("A transaction has already started")
  if self._spanner_transaction_started:
  raise OperationalError(
- "Beginning a new transaction is not allowed when a transaction is already running"
+ "Beginning a new transaction is not allowed when a transaction "
+ "is already running"
  )
  self._transaction_begin_marked = True
 
@@ -430,41 +429,37 @@ def commit(self):
  return
 
  self.run_prior_DDL_statements()
- if self._spanner_transaction_started:
- try:
- if self.read_only:
- self._snapshot = None
- else:
- self._transaction.commit()
-
- self._release_session()
- self._statements = []
- self._transaction_begin_marked = False
- except Aborted:
- self.retry_transaction()
- self.commit()
+ try:
+ if self._spanner_transaction_started and not self._read_only:
+ self._transaction.commit()
+ except Aborted:
+ self.retry_transaction()
+ self.commit()
+ finally:
+ self._release_session()
+ self._statements = []
+ self._transaction_begin_marked = False
+ self._spanner_transaction_started = False
 
  def rollback(self):
  """Rolls back any pending transaction.
 
  This is a no-op if there is no active client transaction.
  """
-
  if not self._client_transaction_started:
  warnings.warn(
  CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
  )
  return
 
- if self._spanner_transaction_started:
- if self.read_only:
- self._snapshot = None
- else:
+ try:
+ if self._spanner_transaction_started and not self._read_only:
  self._transaction.rollback()
-
+ finally:
  self._release_session()
  self._statements = []
  self._transaction_begin_marked = False
+ self._spanner_transaction_started = False
 
  @check_not_closed
  def cursor(self):