feat!: Implementation of batch ddl in dbapi

google/cloud/spanner_dbapi/batch_dml_executor.py → google/cloud/spanner_dbapi/batch_executor.py

@@ -16,6 +16,8 @@
 
 from enum import Enum
 from typing import TYPE_CHECKING, List
+
+from google.cloud.spanner_dbapi import parse_utils
 from google.cloud.spanner_dbapi.parsed_statement import (
     ParsedStatement,
     StatementType,
@@ -28,6 +30,52 @@
 
 if TYPE_CHECKING:
     from google.cloud.spanner_dbapi.cursor import Cursor
+    from google.cloud.spanner_dbapi.connection import Connection
+
+
+class BatchDdlExecutor:
+    """Executor that is used when a DDL batch is started. These batches only
+    accept DDL statements. All DDL statements are buffered locally and sent to
+    Spanner when runBatch() is called.
+
+    :type "Connection": :class:`~google.cloud.spanner_dbapi.connection.Connection`
+    :param connection:
+    """
+
+    def __init__(self, connection: "Connection"):
+        self._connection = connection
+        self._statements: List[str] = []
+
+    def execute_statement(self, parsed_statement: ParsedStatement):
+        """Executes the statement when ddl batch is active by buffering the
+        statement in-memory.
+
+        This method is internal and not for public use as it can change anytime.
+
+        :type parsed_statement: ParsedStatement
+        :param parsed_statement: parsed statement containing sql query
+        """
+        from google.cloud.spanner_dbapi import ProgrammingError
+
+        if parsed_statement.statement_type != StatementType.DDL:
+            raise ProgrammingError("Only DDL statements are allowed in batch DDL mode.")
+        self._statements.extend(
+            parse_utils.parse_and_get_ddl_statements(parsed_statement.statement.sql)
+        )
+
+    def run_batch(self):
+        """Executes all the buffered statements on the active ddl batch by
+        making a call to Spanner.
+
+        This method is internal and not for public use as it can change anytime.
+        """
+        from google.cloud.spanner_dbapi import ProgrammingError
+
+        if self._connection._client_transaction_started:
+            raise ProgrammingError(
+                "Cannot execute DDL statement when transaction is already active."
+            )
+        return self._connection.database.update_ddl(self._statements).result()
 
 
 class BatchDmlExecutor:
@@ -48,10 +96,13 @@ def execute_statement(self, parsed_statement: ParsedStatement):
         """Executes the statement when dml batch is active by buffering the
         statement in-memory.
 
+        This method is internal and not for public use as it can change anytime.
+
         :type parsed_statement: ParsedStatement
         :param parsed_statement: parsed statement containing sql query and query
          params
         """
+
         from google.cloud.spanner_dbapi import ProgrammingError
 
         if (
@@ -61,16 +112,20 @@ def execute_statement(self, parsed_statement: ParsedStatement):
             raise ProgrammingError("Only DML statements are allowed in batch DML mode.")
         self._statements.append(parsed_statement.statement)
 
-    def run_batch_dml(self):
+    def run_batch(self):
         """Executes all the buffered statements on the active dml batch by
         making a call to Spanner.
+
+        This method is internal and not for public use as it can change anytime.
         """
         return run_batch_dml(self._cursor, self._statements)
 
 
 def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
     """Executes all the dml statements by making a batch call to Spanner.
 
+    This method is internal and not for public use as it can change anytime.
+
     :type cursor: Cursor
     :param cursor: Database Cursor object
 

google/cloud/spanner_dbapi/client_side_statement_executor.py

@@ -85,6 +85,9 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement):
             TypeCode.TIMESTAMP,
             column_values,
         )
+    if statement_type == ClientSideStatementType.START_BATCH_DDL:
+        connection.start_batch_ddl()
+        return None
     if statement_type == ClientSideStatementType.START_BATCH_DML:
         connection.start_batch_dml(cursor)
         return None

google/cloud/spanner_dbapi/client_side_statement_parser.py

@@ -30,6 +30,7 @@
 RE_SHOW_READ_TIMESTAMP = re.compile(
     r"^\s*(SHOW)\s+(VARIABLE)\s+(READ_TIMESTAMP)", re.IGNORECASE
 )
+RE_START_BATCH_DDL = re.compile(r"^\s*(START)\s+(BATCH)\s+(DDL)", re.IGNORECASE)
 RE_START_BATCH_DML = re.compile(r"^\s*(START)\s+(BATCH)\s+(DML)", re.IGNORECASE)
 RE_RUN_BATCH = re.compile(r"^\s*(RUN)\s+(BATCH)", re.IGNORECASE)
 RE_ABORT_BATCH = re.compile(r"^\s*(ABORT)\s+(BATCH)", re.IGNORECASE)
@@ -62,6 +63,8 @@ def parse_stmt(query):
         client_side_statement_type = ClientSideStatementType.SHOW_COMMIT_TIMESTAMP
     elif RE_SHOW_READ_TIMESTAMP.match(query):
         client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP
+    elif RE_START_BATCH_DDL.match(query):
+        client_side_statement_type = ClientSideStatementType.START_BATCH_DDL
     elif RE_START_BATCH_DML.match(query):
         client_side_statement_type = ClientSideStatementType.START_BATCH_DML
     elif RE_BEGIN.match(query):

google/cloud/spanner_dbapi/connection.py

@@ -19,7 +19,11 @@
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.cloud import spanner_v1 as spanner
 from google.cloud.spanner_dbapi import partition_helper
-from google.cloud.spanner_dbapi.batch_dml_executor import BatchMode, BatchDmlExecutor
+from google.cloud.spanner_dbapi.batch_executor import (
+    BatchMode,
+    BatchDmlExecutor,
+    BatchDdlExecutor,
+)
 from google.cloud.spanner_dbapi.parse_utils import _get_statement_type
 from google.cloud.spanner_dbapi.parsed_statement import (
     StatementType,
@@ -91,7 +95,9 @@ class Connection:
         should end a that a new one should be started when the next statement is executed.
     """
 
-    def __init__(self, instance, database=None, read_only=False):
+    def __init__(
+        self, instance, database=None, read_only=False, buffer_ddl_statements=False
+    ):
         self._instance = instance
         self._database = database
         self._ddl_statements = []
@@ -114,8 +120,10 @@ def __init__(self, instance, database=None, read_only=False):
         # made atleast one call to Spanner.
         self._spanner_transaction_started = False
         self._batch_mode = BatchMode.NONE
+        self._batch_ddl_executor: BatchDdlExecutor = None
         self._batch_dml_executor: BatchDmlExecutor = None
         self._transaction_helper = TransactionRetryHelper(self)
+        self._buffer_ddl_statements = buffer_ddl_statements
 
     @property
     def autocommit(self):
@@ -126,6 +134,25 @@ def autocommit(self):
         """
         return self._autocommit
 
+    @property
+    def buffer_ddl_statements(self):
+        """Whether to buffer ddl statements for this connection.
+        If the connection is in auto commit mode then this flag doesn't have any
+        significance as ddl statements would be executed as they come.
+
+        For connection not in auto commit mode:
+        If enabled ddl statements would be buffered at client and not executed
+        at cloud spanner. When a non ddl statement comes or a transaction is
+        committed then all the existing buffered ddl statements would be executed.
+
+        If disabled then its an error to execute ddl statement in autocommit mode.
+
+
+        :rtype: bool
+        :returns: _buffer_ddl_statements flag value.
+        """
+        return self._buffer_ddl_statements
+
     @autocommit.setter
     def autocommit(self, value):
         """Change this connection autocommit mode. Setting this value to True
@@ -365,7 +392,8 @@ def commit(self):
             )
             return
 
-        self.run_prior_DDL_statements()
+        if self.buffer_ddl_statements:
+            self.run_prior_DDL_statements()
         try:
             if self._spanner_transaction_started and not self._read_only:
                 self._transaction.commit()
@@ -463,8 +491,42 @@ def validate(self):
                     "Expected: [[1]]" % result
                 )
 
+    @check_not_closed
+    def start_batch_ddl(self):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
+        if self._batch_mode is not BatchMode.NONE:
+            raise ProgrammingError(
+                "Cannot start a DDL batch when a batch is already active"
+            )
+        if self.read_only:
+            raise ProgrammingError(
+                "Cannot start a DDL batch when the connection is in read-only mode"
+            )
+        if self.buffer_ddl_statements:
+            raise ProgrammingError(
+                "Cannot start a DDL batch when _buffer_ddl_statements flag is True"
+            )
+        self._batch_mode = BatchMode.DDL
+        self._batch_ddl_executor = BatchDdlExecutor(self)
+
+    @check_not_closed
+    def execute_batch_ddl_statement(self, parsed_statement: ParsedStatement):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
+        if self._batch_mode is not BatchMode.DDL:
+            raise ProgrammingError(
+                "Cannot execute statement when the BatchMode is not DDL"
+            )
+        self._batch_ddl_executor.execute_statement(parsed_statement)
+
     @check_not_closed
     def start_batch_dml(self, cursor):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
         if self._batch_mode is not BatchMode.NONE:
             raise ProgrammingError(
                 "Cannot start a DML batch when a batch is already active"
@@ -478,6 +540,9 @@ def start_batch_dml(self, cursor):
 
     @check_not_closed
     def execute_batch_dml_statement(self, parsed_statement: ParsedStatement):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
         if self._batch_mode is not BatchMode.DML:
             raise ProgrammingError(
                 "Cannot execute statement when the BatchMode is not DML"
@@ -486,22 +551,34 @@ def execute_batch_dml_statement(self, parsed_statement: ParsedStatement):
 
     @check_not_closed
     def run_batch(self):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
+        result_set = None
         if self._batch_mode is BatchMode.NONE:
             raise ProgrammingError("Cannot run a batch when the BatchMode is not set")
         try:
             if self._batch_mode is BatchMode.DML:
-                many_result_set = self._batch_dml_executor.run_batch_dml()
+                result_set = self._batch_dml_executor.run_batch()
+            elif self._batch_mode is BatchMode.DDL:
+                self._batch_ddl_executor.run_batch()
         finally:
             self._batch_mode = BatchMode.NONE
             self._batch_dml_executor = None
-        return many_result_set
+            self._batch_ddl_executor = None
+        return result_set
 
     @check_not_closed
     def abort_batch(self):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
         if self._batch_mode is BatchMode.NONE:
             raise ProgrammingError("Cannot abort a batch when the BatchMode is not set")
         if self._batch_mode is BatchMode.DML:
             self._batch_dml_executor = None
+        if self._batch_mode is BatchMode.DDL:
+            self._batch_ddl_executor = None
         self._batch_mode = BatchMode.NONE
 
     @check_not_closed
@@ -510,6 +587,9 @@ def partition_query(
         parsed_statement: ParsedStatement,
         query_options=None,
     ):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
         statement = parsed_statement.statement
         partitioned_query = parsed_statement.client_side_statement_params[0]
         self._partitioned_query_validation(partitioned_query, statement)
@@ -534,6 +614,9 @@ def partition_query(
 
     @check_not_closed
     def run_partition(self, encoded_partition_id):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
         partition_id: PartitionId = partition_helper.decode_from_string(
             encoded_partition_id
         )
@@ -550,6 +633,9 @@ def run_partitioned_query(
         self,
         parsed_statement: ParsedStatement,
     ):
+        """
+        This method is internal and not for public use as it can change anytime.
+        """
         statement = parsed_statement.statement
         partitioned_query = parsed_statement.client_side_statement_params[0]
         self._partitioned_query_validation(partitioned_query, statement)
@@ -584,10 +670,23 @@ def connect(
     pool=None,
     user_agent=None,
     client=None,
+    buffer_ddl_statements=False,
     route_to_leader_enabled=True,
 ):
     """Creates a connection to a Google Cloud Spanner database.
 
+    :type buffer_ddl_statements: bool
+    :param buffer_ddl_statements: Whether to buffer ddl statements at client
+        side. If the connection is in auto commit mode  then this flag doesn't
+        have any significance as ddl statements would be executed as they come.
+
+        For connection not in auto commit mode:
+        If enabled ddl statements would be buffered at client and not executed
+        at cloud spanner. When a non ddl statement comes or a transaction is
+        committed then all the existing buffered ddl statements would be executed.
+
+        If disabled then its an error to execute ddl statement in autocommit mode.
+
     :type instance_id: str
     :param instance_id: The ID of the instance to connect to.
 
@@ -658,7 +757,9 @@ def connect(
 
     instance = client.instance(instance_id)
     conn = Connection(
-        instance, instance.database(database_id, pool=pool) if database_id else None
+        instance,
+        instance.database(database_id, pool=pool) if database_id else None,
+        buffer_ddl_statements=buffer_ddl_statements,
     )
     if pool is not None:
         conn._own_pool = False