feat(spanner): add DataBoost and auto_partition_mode support to DBAPI driver

packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_executor.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 from typing import TYPE_CHECKING, Union
 
+from google.cloud.spanner_dbapi.exceptions import ProgrammingError
 from google.cloud.spanner_v1 import TransactionOptions
 
 if TYPE_CHECKING:
-    from google.cloud.spanner_dbapi import ProgrammingError
     from google.cloud.spanner_dbapi.cursor import Cursor
 
 from google.cloud.spanner_dbapi.parsed_statement import (
@@ -108,6 +108,45 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement):
         return connection.run_partitioned_query(parsed_statement)
     if statement_type == ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE:
         return connection._set_autocommit_dml_mode(parsed_statement)
+    if statement_type == ClientSideStatementType.SET_DATA_BOOST_ENABLED:
+        connection.data_boost_enabled = _parse_bool(
+            parsed_statement.client_side_statement_params[0],
+            "DATA_BOOST_ENABLED",
+        )
+        return None
+    if statement_type == ClientSideStatementType.SHOW_DATA_BOOST_ENABLED:
+        column_values.append(connection.data_boost_enabled)
+        return _get_streamed_result_set(
+            "DATA_BOOST_ENABLED",
+            TypeCode.BOOL,
+            column_values,
+        )
+    if statement_type == ClientSideStatementType.SET_AUTO_PARTITION_MODE:
+        connection.auto_partition_mode = _parse_bool(
+            parsed_statement.client_side_statement_params[0],
+            "AUTO_PARTITION_MODE",
+        )
+        return None
+    if statement_type == ClientSideStatementType.SHOW_AUTO_PARTITION_MODE:
+        column_values.append(connection.auto_partition_mode)
+        return _get_streamed_result_set(
+            "AUTO_PARTITION_MODE",
+            TypeCode.BOOL,
+            column_values,
+        )
+    return None
+
+
+_BOOL_MAP = {"true": True, "false": False}
+
+
+def _parse_bool(raw_val: str, var_name: str) -> bool:
+    cleaned = raw_val.strip().rstrip(";").strip().strip("'\"").lower()
+    if cleaned not in _BOOL_MAP:
+        raise ProgrammingError(
+            f"Invalid value for {var_name}: '{raw_val}'. Expected TRUE or FALSE."
+        )
+    return _BOOL_MAP[cleaned]
 
 
 def _get_streamed_result_set(column_name, type_code, column_values):

packages/google-cloud-spanner/google/cloud/spanner_dbapi/client_side_statement_parser.py

@@ -44,6 +44,18 @@
 RE_SET_AUTOCOMMIT_DML_MODE = re.compile(
     r"^\s*(SET)\s+(AUTOCOMMIT_DML_MODE)\s+(=)\s+(.+)", re.IGNORECASE
 )
+RE_SET_DATA_BOOST_ENABLED = re.compile(
+    r"^\s*(SET)\s+(DATA_BOOST_ENABLED)\s+(=)\s+(.+)", re.IGNORECASE
+)
+RE_SHOW_DATA_BOOST_ENABLED = re.compile(
+    r"^\s*(SHOW)\s+(VARIABLE)\s+(DATA_BOOST_ENABLED)\s*;?\s*$", re.IGNORECASE
+)
+RE_SET_AUTO_PARTITION_MODE = re.compile(
+    r"^\s*(SET)\s+(AUTO_PARTITION_MODE)\s+(=)\s+(.+)", re.IGNORECASE
+)
+RE_SHOW_AUTO_PARTITION_MODE = re.compile(
+    r"^\s*(SHOW)\s+(VARIABLE)\s+(AUTO_PARTITION_MODE)\s*;?\s*$", re.IGNORECASE
+)
 
 
 def parse_stmt(query):
@@ -68,6 +80,10 @@ 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_SHOW_DATA_BOOST_ENABLED.match(query):
+        client_side_statement_type = ClientSideStatementType.SHOW_DATA_BOOST_ENABLED
+    elif RE_SHOW_AUTO_PARTITION_MODE.match(query):
+        client_side_statement_type = ClientSideStatementType.SHOW_AUTO_PARTITION_MODE
     elif RE_START_BATCH_DML.match(query):
         client_side_statement_type = ClientSideStatementType.START_BATCH_DML
     elif RE_BEGIN.match(query):
@@ -96,6 +112,14 @@ def parse_stmt(query):
         match = re.search(RE_SET_AUTOCOMMIT_DML_MODE, query)
         client_side_statement_params.append(match.group(4))
         client_side_statement_type = ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE
+    elif RE_SET_DATA_BOOST_ENABLED.match(query):
+        match = re.search(RE_SET_DATA_BOOST_ENABLED, query)
+        client_side_statement_params.append(match.group(4))
+        client_side_statement_type = ClientSideStatementType.SET_DATA_BOOST_ENABLED
+    elif RE_SET_AUTO_PARTITION_MODE.match(query):
+        match = re.search(RE_SET_AUTO_PARTITION_MODE, query)
+        client_side_statement_params.append(match.group(4))
+        client_side_statement_type = ClientSideStatementType.SET_AUTO_PARTITION_MODE
     if client_side_statement_type is not None:
         return ParsedStatement(
             StatementType.CLIENT_SIDE,

packages/google-cloud-spanner/google/cloud/spanner_dbapi/connection.py

@@ -20,7 +20,6 @@
 from google.api_core.exceptions import Aborted
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.auth.credentials import AnonymousCredentials
-
 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 BatchDmlExecutor, BatchMode
@@ -91,10 +90,29 @@ class Connection:
         the read-only transaction is semantically the same, and only indicates that the read-only transaction
         should end a that a new one should be started when the next statement is executed.
 
+    :type data_boost_enabled: bool
+    :param data_boost_enabled: (Optional) Whether to enable DataBoost for
+        partitioned queries executed via this connection. Defaults to False.
+        Note that DataBoost is only supported for partitioned query execution.
+
+    :type auto_partition_mode: bool
+    :param auto_partition_mode: (Optional) Whether to enable auto partition mode
+        for queries executed via this connection. When True, queries on read-only
+        or autocommit connections are automatically partitioned and executed in parallel.
+        Defaults to False.
+
     **kwargs: Initial value for connection variables.
     """
 
-    def __init__(self, instance, database=None, read_only=False, **kwargs):
+    def __init__(
+        self,
+        instance,
+        database=None,
+        read_only=False,
+        data_boost_enabled=False,
+        auto_partition_mode=False,
+        **kwargs,
+    ):
         self._instance = instance
         self._database = database
         self._ddl_statements = []
@@ -110,6 +128,8 @@ def __init__(self, instance, database=None, read_only=False, **kwargs):
         # connection close
         self._own_pool = True
         self._read_only = read_only
+        self._data_boost_enabled = bool(data_boost_enabled)
+        self._auto_partition_mode = bool(auto_partition_mode)
         self._staleness = None
         self.request_priority = None
         self._transaction_begin_marked = False
@@ -123,6 +143,47 @@ def __init__(self, instance, database=None, read_only=False, **kwargs):
         self._autocommit_dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
         self._connection_variables = kwargs
 
+    @property
+    def data_boost_enabled(self):
+        """Flag: whether DataBoost is enabled for partitioned queries on this connection.
+
+        Note that DataBoost is only supported for partitioned query execution.
+
+        Returns:
+            bool: True if DataBoost is enabled, False otherwise.
+        """
+        return self._data_boost_enabled
+
+    @data_boost_enabled.setter
+    def data_boost_enabled(self, value):
+        """Change the DataBoost enablement state for partitioned queries on this connection.
+
+        :type value: bool
+        :param value: New data_boost_enabled state.
+        """
+        self._data_boost_enabled = bool(value)
+
+    @property
+    def auto_partition_mode(self):
+        """Flag: whether auto partition mode is enabled for queries on this connection.
+
+        When enabled, standard queries executed on read-only or autocommit connections
+        are automatically partitioned and executed in parallel via run_partitioned_query.
+
+        Returns:
+            bool: True if auto partition mode is enabled, False otherwise.
+        """
+        return self._auto_partition_mode
+
+    @auto_partition_mode.setter
+    def auto_partition_mode(self, value):
+        """Change the auto partition mode enablement state for this connection.
+
+        :type value: bool
+        :param value: New auto_partition_mode state.
+        """
+        self._auto_partition_mode = bool(value)
+
     @property
     def spanner_client(self):
         """Client for interacting with Cloud Spanner API. This property exposes
@@ -638,11 +699,15 @@ def partition_query(
         self,
         parsed_statement: ParsedStatement,
         query_options=None,
+        data_boost_enabled=None,
     ):
         statement = parsed_statement.statement
         partitioned_query = parsed_statement.client_side_statement_params[0]
         self._partitioned_query_validation(partitioned_query, statement)
 
+        if data_boost_enabled is None:
+            data_boost_enabled = self.data_boost_enabled
+
         batch_snapshot = self._database.batch_snapshot()
         partition_ids = []
         partitions = list(
@@ -651,6 +716,7 @@ def partition_query(
                 statement.params,
                 statement.param_types,
                 query_options=query_options,
+                data_boost_enabled=data_boost_enabled,
             )
         )
 
@@ -684,7 +750,10 @@ def run_partitioned_query(
         self._partitioned_query_validation(partitioned_query, statement)
         batch_snapshot = self._database.batch_snapshot()
         return batch_snapshot.run_partitioned_query(
-            partitioned_query, statement.params, statement.param_types
+            partitioned_query,
+            statement.params,
+            statement.param_types,
+            data_boost_enabled=self.data_boost_enabled,
         )
 
     @check_not_closed
@@ -748,6 +817,8 @@ def connect(
     client_certificate=None,
     client_key=None,
     instance_type=None,
+    data_boost_enabled=False,
+    auto_partition_mode=False,
     **kwargs,
 ):
     """Creates a connection to a Google Cloud Spanner database.
@@ -795,6 +866,17 @@ def connect(
     :param database_role: (Optional) The database role to connect as when using
         fine-grained access controls.
 
+    :type data_boost_enabled: bool
+    :param data_boost_enabled: (Optional) Whether to enable DataBoost for
+        partitioned queries executed via this connection. Defaults to False.
+        Note that DataBoost is only supported for partitioned query execution.
+
+    :type auto_partition_mode: bool
+    :param auto_partition_mode: (Optional) Whether to enable auto partition mode
+        for queries executed via this connection. When True, queries on read-only
+        or autocommit connections are automatically partitioned and executed in parallel.
+        Defaults to False.
+
     **kwargs: Initial value for connection variables.
 
 
@@ -909,7 +991,13 @@ def connect(
         database = instance.database(
             database_id, pool=pool, database_role=database_role, logger=logger
         )
-    conn = Connection(instance, database, **kwargs)
+    conn = Connection(
+        instance,
+        database,
+        data_boost_enabled=data_boost_enabled,
+        auto_partition_mode=auto_partition_mode,
+        **kwargs,
+    )
     if pool is not None:
         conn._own_pool = False
 

packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py

@@ -25,7 +25,6 @@
     InvalidArgument,
     OutOfRange,
 )
-
 from google.cloud import spanner_v1 as spanner
 from google.cloud.spanner_dbapi import (
     _helpers,
@@ -308,6 +307,11 @@ def _execute(self, sql, args=None, call_from_execute_many=False):
                         self._itr = PeekIterator(self._result_set)
             elif self.connection._batch_mode == BatchMode.DML:
                 self.connection.execute_batch_dml_statement(self._parsed_statement)
+            elif (
+                self._parsed_statement.statement_type == StatementType.QUERY
+                and self.connection.auto_partition_mode
+            ):
+                self._handle_auto_partition_query(sql, args or None)
             elif self.connection.read_only or (
                 not self.connection._client_transaction_started
                 and self._parsed_statement.statement_type == StatementType.QUERY
@@ -581,6 +585,27 @@ def _handle_DQL(self, sql, params):
                 self.connection._transaction = None
                 self._handle_DQL_with_snapshot(snapshot, sql, params)
 
+    def _handle_auto_partition_query(self, sql, params):
+        if self.connection.database is None:
+            raise ValueError("Database needs to be passed for this operation")
+        if (
+            not self.connection.read_only
+            and self.connection._client_transaction_started
+        ):
+            raise ProgrammingError(
+                "Partitioned query is not supported, because the connection is in a read/write transaction."
+            )
+        sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params)
+        batch_snapshot = self.connection.database.batch_snapshot()
+        self._result_set = batch_snapshot.run_partitioned_query(
+            sql,
+            params=params,
+            param_types=get_param_types(params),
+            data_boost_enabled=self.connection.data_boost_enabled,
+        )
+        self._itr = self._result_set
+        self._row_count = None
+
     def __enter__(self):
         return self
 

packages/google-cloud-spanner/google/cloud/spanner_dbapi/parsed_statement.py

@@ -38,6 +38,10 @@ class ClientSideStatementType(Enum):
     RUN_PARTITION = 10
     RUN_PARTITIONED_QUERY = 11
     SET_AUTOCOMMIT_DML_MODE = 12
+    SET_DATA_BOOST_ENABLED = 13
+    SHOW_DATA_BOOST_ENABLED = 14
+    SET_AUTO_PARTITION_MODE = 15
+    SHOW_AUTO_PARTITION_MODE = 16
 
 
 class AutocommitDmlMode(Enum):