google-cloud-spanner 3.41.0__tar.gz → 3.43.0__tar.gz

{google-cloud-spanner-3.41.0/google_cloud_spanner.egg-info → google-cloud-spanner-3.43.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: google-cloud-spanner
-Version: 3.41.0
+Version: 3.43.0
 Summary: Google Cloud Spanner API client library
 Home-page: https://github.com/googleapis/python-spanner
 Author: Google LLC
@@ -30,6 +30,7 @@ Requires-Dist: sqlparse>=0.4.4
 Requires-Dist: proto-plus<2.0.0dev,>=1.22.2; python_version >= "3.11"
 Requires-Dist: protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5
 Requires-Dist: deprecated>=1.2.14
+Requires-Dist: grpc-interceptor>=0.15.4
 Provides-Extra: tracing
 Requires-Dist: opentelemetry-api>=1.1.0; extra == "tracing"
 Requires-Dist: opentelemetry-sdk>=1.1.0; extra == "tracing"

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_admin_database_v1/gapic_version.py

@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "3.41.0"  # {x-release-please-version}
+__version__ = "3.43.0"  # {x-release-please-version}

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_admin_database_v1/types/spanner_database_admin.py

@@ -355,6 +355,26 @@ class CreateDatabaseRequest(proto.Message):
         database_dialect (google.cloud.spanner_admin_database_v1.types.DatabaseDialect):
             Optional. The dialect of the Cloud Spanner
             Database.
+        proto_descriptors (bytes):
+            Optional. Proto descriptors used by CREATE/ALTER PROTO
+            BUNDLE statements in 'extra_statements' above. Contains a
+            protobuf-serialized
+            `google.protobuf.FileDescriptorSet <https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/descriptor.proto>`__.
+            To generate it,
+            `install <https://grpc.io/docs/protoc-installation/>`__ and
+            run ``protoc`` with --include_imports and
+            --descriptor_set_out. For example, to generate for
+            moon/shot/app.proto, run
+
+            ::
+
+               $protoc  --proto_path=/app_path --proto_path=/lib_path \
+                        --include_imports \
+                        --descriptor_set_out=descriptors.data \
+                        moon/shot/app.proto
+
+            For more details, see protobuffer `self
+            description <https://developers.google.com/protocol-buffers/docs/techniques#self-description>`__.
     """
 
     parent: str = proto.Field(
@@ -379,6 +399,10 @@ class CreateDatabaseRequest(proto.Message):
         number=5,
         enum=common.DatabaseDialect,
     )
+    proto_descriptors: bytes = proto.Field(
+        proto.BYTES,
+        number=6,
+    )
 
 
 class CreateDatabaseMetadata(proto.Message):
@@ -521,6 +545,25 @@ class UpdateDatabaseDdlRequest(proto.Message):
             underscore. If the named operation already exists,
             [UpdateDatabaseDdl][google.spanner.admin.database.v1.DatabaseAdmin.UpdateDatabaseDdl]
             returns ``ALREADY_EXISTS``.
+        proto_descriptors (bytes):
+            Optional. Proto descriptors used by CREATE/ALTER PROTO
+            BUNDLE statements. Contains a protobuf-serialized
+            `google.protobuf.FileDescriptorSet <https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/descriptor.proto>`__.
+            To generate it,
+            `install <https://grpc.io/docs/protoc-installation/>`__ and
+            run ``protoc`` with --include_imports and
+            --descriptor_set_out. For example, to generate for
+            moon/shot/app.proto, run
+
+            ::
+
+               $protoc  --proto_path=/app_path --proto_path=/lib_path \
+                        --include_imports \
+                        --descriptor_set_out=descriptors.data \
+                        moon/shot/app.proto
+
+            For more details, see protobuffer `self
+            description <https://developers.google.com/protocol-buffers/docs/techniques#self-description>`__.
     """
 
     database: str = proto.Field(
@@ -535,6 +578,10 @@ class UpdateDatabaseDdlRequest(proto.Message):
         proto.STRING,
         number=3,
     )
+    proto_descriptors: bytes = proto.Field(
+        proto.BYTES,
+        number=4,
+    )
 
 
 class DdlStatementActionInfo(proto.Message):
@@ -682,12 +729,22 @@ class GetDatabaseDdlResponse(proto.Message):
             A list of formatted DDL statements defining
             the schema of the database specified in the
             request.
+        proto_descriptors (bytes):
+            Proto descriptors stored in the database. Contains a
+            protobuf-serialized
+            `google.protobuf.FileDescriptorSet <https://github.com/protocolbuffers/protobuf/blob/main/src/google/protobuf/descriptor.proto>`__.
+            For more details, see protobuffer `self
+            description <https://developers.google.com/protocol-buffers/docs/techniques#self-description>`__.
     """
 
     statements: MutableSequence[str] = proto.RepeatedField(
         proto.STRING,
         number=1,
     )
+    proto_descriptors: bytes = proto.Field(
+        proto.BYTES,
+        number=2,
+    )
 
 
 class ListDatabaseOperationsRequest(proto.Message):

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_admin_instance_v1/gapic_version.py

@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "3.41.0"  # {x-release-please-version}
+__version__ = "3.43.0"  # {x-release-please-version}

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_dbapi/_helpers.py

@@ -18,13 +18,13 @@ from google.cloud.spanner_v1 import param_types
 SQL_LIST_TABLES = """
 SELECT table_name
 FROM information_schema.tables
-WHERE table_catalog = '' AND table_schema = ''
+WHERE table_catalog = '' AND table_schema = @table_schema
 """
 
 SQL_GET_TABLE_COLUMN_SCHEMA = """
 SELECT COLUMN_NAME, IS_NULLABLE, SPANNER_TYPE
 FROM INFORMATION_SCHEMA.COLUMNS
-WHERE TABLE_SCHEMA = '' AND TABLE_NAME = @table_name
+WHERE TABLE_SCHEMA = @table_schema AND TABLE_NAME = @table_name
 """
 
 # This table maps spanner_types to Spanner's data type sizes as per

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_dbapi/batch_dml_executor.py

@@ -16,7 +16,6 @@ from __future__ import annotations
 
 from enum import Enum
 from typing import TYPE_CHECKING, List
-from google.cloud.spanner_dbapi.checksum import ResultsChecksum
 from google.cloud.spanner_dbapi.parsed_statement import (
     ParsedStatement,
     StatementType,
@@ -80,8 +79,10 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
     """
     from google.cloud.spanner_dbapi import OperationalError
 
-    connection = cursor.connection
     many_result_set = StreamedManyResultSets()
+    if not statements:
+        return many_result_set
+    connection = cursor.connection
     statements_tuple = []
     for statement in statements:
         statements_tuple.append(statement.get_tuple())
@@ -90,28 +91,26 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
         many_result_set.add_iter(res)
         cursor._row_count = sum([max(val, 0) for val in res])
     else:
-        retried = False
         while True:
             try:
                 transaction = connection.transaction_checkout()
                 status, res = transaction.batch_update(statements_tuple)
-                many_result_set.add_iter(res)
-                res_checksum = ResultsChecksum()
-                res_checksum.consume_result(res)
-                res_checksum.consume_result(status.code)
-                if not retried:
-                    connection._statements.append((statements, res_checksum))
-                cursor._row_count = sum([max(val, 0) for val in res])
-
                 if status.code == ABORTED:
                     connection._transaction = None
                     raise Aborted(status.message)
                 elif status.code != OK:
                     raise OperationalError(status.message)
+
+                cursor._batch_dml_rows_count = res
+                many_result_set.add_iter(res)
+                cursor._row_count = sum([max(val, 0) for val in res])
                 return many_result_set
             except Aborted:
-                connection.retry_transaction()
-                retried = True
+                # We are raising it so it could be handled in transaction_helper.py and is retried
+                if cursor._in_retry_mode:
+                    raise
+                else:
+                    connection._transaction_helper.retry_transaction()
 
 
 def _do_batch_update(transaction, statements):

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_dbapi/checksum.py

@@ -62,6 +62,8 @@ class ResultsChecksum:
 
 
 def _compare_checksums(original, retried):
+    from google.cloud.spanner_dbapi.transaction_helper import RETRY_ABORTED_ERROR
+
     """Compare the given checksums.
 
     Raise an error if the given checksums are not equal.
@@ -75,6 +77,4 @@ def _compare_checksums(original, retried):
     :raises: :exc:`google.cloud.spanner_dbapi.exceptions.RetryAborted` in case if checksums are not equal.
     """
     if retried != original:
-        raise RetryAborted(
-            "The transaction was aborted and could not be retried due to a concurrent modification."
-        )
+        raise RetryAborted(RETRY_ABORTED_ERROR)

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_dbapi/client_side_statement_executor.py

@@ -103,6 +103,10 @@ def execute(cursor: "Cursor", parsed_statement: ParsedStatement):
         return connection.run_partition(
             parsed_statement.client_side_statement_params[0]
         )
+    if statement_type == ClientSideStatementType.RUN_PARTITIONED_QUERY:
+        return connection.run_partitioned_query(parsed_statement)
+    if statement_type == ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE:
+        return connection._set_autocommit_dml_mode(parsed_statement)
 
 
 def _get_streamed_result_set(column_name, type_code, column_values):

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_dbapi/client_side_statement_parser.py

@@ -35,6 +35,12 @@ RE_RUN_BATCH = re.compile(r"^\s*(RUN)\s+(BATCH)", re.IGNORECASE)
 RE_ABORT_BATCH = re.compile(r"^\s*(ABORT)\s+(BATCH)", re.IGNORECASE)
 RE_PARTITION_QUERY = re.compile(r"^\s*(PARTITION)\s+(.+)", re.IGNORECASE)
 RE_RUN_PARTITION = re.compile(r"^\s*(RUN)\s+(PARTITION)\s+(.+)", re.IGNORECASE)
+RE_RUN_PARTITIONED_QUERY = re.compile(
+    r"^\s*(RUN)\s+(PARTITIONED)\s+(QUERY)\s+(.+)", re.IGNORECASE
+)
+RE_SET_AUTOCOMMIT_DML_MODE = re.compile(
+    r"^\s*(SET)\s+(AUTOCOMMIT_DML_MODE)\s+(=)\s+(.+)", re.IGNORECASE
+)
 
 
 def parse_stmt(query):
@@ -53,28 +59,36 @@ def parse_stmt(query):
     client_side_statement_params = []
     if RE_COMMIT.match(query):
         client_side_statement_type = ClientSideStatementType.COMMIT
-    if RE_BEGIN.match(query):
-        client_side_statement_type = ClientSideStatementType.BEGIN
-    if RE_ROLLBACK.match(query):
+    elif RE_ROLLBACK.match(query):
         client_side_statement_type = ClientSideStatementType.ROLLBACK
-    if RE_SHOW_COMMIT_TIMESTAMP.match(query):
+    elif RE_SHOW_COMMIT_TIMESTAMP.match(query):
         client_side_statement_type = ClientSideStatementType.SHOW_COMMIT_TIMESTAMP
-    if RE_SHOW_READ_TIMESTAMP.match(query):
+    elif RE_SHOW_READ_TIMESTAMP.match(query):
         client_side_statement_type = ClientSideStatementType.SHOW_READ_TIMESTAMP
-    if RE_START_BATCH_DML.match(query):
+    elif RE_START_BATCH_DML.match(query):
         client_side_statement_type = ClientSideStatementType.START_BATCH_DML
-    if RE_RUN_BATCH.match(query):
+    elif RE_BEGIN.match(query):
+        client_side_statement_type = ClientSideStatementType.BEGIN
+    elif RE_RUN_BATCH.match(query):
         client_side_statement_type = ClientSideStatementType.RUN_BATCH
-    if RE_ABORT_BATCH.match(query):
+    elif RE_ABORT_BATCH.match(query):
         client_side_statement_type = ClientSideStatementType.ABORT_BATCH
-    if RE_PARTITION_QUERY.match(query):
+    elif RE_RUN_PARTITIONED_QUERY.match(query):
+        match = re.search(RE_RUN_PARTITIONED_QUERY, query)
+        client_side_statement_params.append(match.group(4))
+        client_side_statement_type = ClientSideStatementType.RUN_PARTITIONED_QUERY
+    elif RE_PARTITION_QUERY.match(query):
         match = re.search(RE_PARTITION_QUERY, query)
         client_side_statement_params.append(match.group(2))
         client_side_statement_type = ClientSideStatementType.PARTITION_QUERY
-    if RE_RUN_PARTITION.match(query):
+    elif RE_RUN_PARTITION.match(query):
         match = re.search(RE_RUN_PARTITION, query)
         client_side_statement_params.append(match.group(3))
         client_side_statement_type = ClientSideStatementType.RUN_PARTITION
+    elif RE_SET_AUTOCOMMIT_DML_MODE.match(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
     if client_side_statement_type is not None:
         return ParsedStatement(
             StatementType.CLIENT_SIDE,

{google-cloud-spanner-3.41.0 → google-cloud-spanner-3.43.0}/google/cloud/spanner_dbapi/connection.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """DB-API Connection for the Google Cloud Spanner."""
-import time
 import warnings
 
 from google.api_core.exceptions import Aborted
@@ -23,19 +22,17 @@ from google.cloud.spanner_dbapi import partition_helper
 from google.cloud.spanner_dbapi.batch_dml_executor import BatchMode, BatchDmlExecutor
 from google.cloud.spanner_dbapi.parse_utils import _get_statement_type
 from google.cloud.spanner_dbapi.parsed_statement import (
-    ParsedStatement,
-    Statement,
     StatementType,
+    AutocommitDmlMode,
 )
 from google.cloud.spanner_dbapi.partition_helper import PartitionId
+from google.cloud.spanner_dbapi.parsed_statement import ParsedStatement, Statement
+from google.cloud.spanner_dbapi.transaction_helper import TransactionRetryHelper
+from google.cloud.spanner_dbapi.cursor import Cursor
 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
-from google.cloud.spanner_dbapi.cursor import Cursor
 from google.cloud.spanner_dbapi.exceptions import (
     InterfaceError,
     OperationalError,
@@ -44,13 +41,10 @@ from google.cloud.spanner_dbapi.exceptions import (
 from google.cloud.spanner_dbapi.version import DEFAULT_USER_AGENT
 from google.cloud.spanner_dbapi.version import PY_VERSION
 
-from google.rpc.code_pb2 import ABORTED
-
 
 CLIENT_TRANSACTION_NOT_STARTED_WARNING = (
     "This method is non-operational as a transaction has not been started."
 )
-MAX_INTERNAL_RETRIES = 50
 
 
 def check_not_closed(function):
@@ -106,9 +100,6 @@ class Connection:
         self._transaction = None
         self._session = None
         self._snapshot = None
-        # SQL statements, which were executed
-        # within the current transaction
-        self._statements = []
 
         self.is_closed = False
         self._autocommit = False
@@ -125,6 +116,27 @@ class Connection:
         self._spanner_transaction_started = False
         self._batch_mode = BatchMode.NONE
         self._batch_dml_executor: BatchDmlExecutor = None
+        self._transaction_helper = TransactionRetryHelper(self)
+        self._autocommit_dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
+
+    @property
+    def spanner_client(self):
+        """Client for interacting with Cloud Spanner API. This property exposes
+        the spanner client so that underlying methods can be accessed.
+        """
+        return self._instance._client
+
+    @property
+    def current_schema(self):
+        """schema name for this connection.
+
+        :rtype: str
+        :returns: the current default schema of this connection. Currently, this
+         is always "" for GoogleSQL and "public" for PostgreSQL databases.
+        """
+        if self.database is None:
+            raise ValueError("database property not set on the connection")
+        return self.database.default_schema_name
 
     @property
     def autocommit(self):
@@ -157,6 +169,23 @@ class Connection:
         """
         return self._database
 
+    @property
+    def autocommit_dml_mode(self):
+        """Modes for executing DML statements in autocommit mode for this connection.
+
+        The DML autocommit modes are:
+        1) TRANSACTIONAL - DML statements are executed as single read-write transaction.
+        After successful execution, the DML statement is guaranteed to have been applied
+        exactly once to the database.
+
+        2) PARTITIONED_NON_ATOMIC - DML statements are executed as partitioned DML transactions.
+        If an error occurs during the execution of the DML statement, it is possible that the
+        statement has been applied to some but not all of the rows specified in the statement.
+
+        :rtype: :class:`~google.cloud.spanner_dbapi.parsed_statement.AutocommitDmlMode`
+        """
+        return self._autocommit_dml_mode
+
     @property
     @deprecated(
         reason="This method is deprecated. Use _spanner_transaction_started field"
@@ -282,82 +311,13 @@ class Connection:
 
         The session will be returned into the sessions pool.
         """
+        if self._session is None:
+            return
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
-        if self._session is not None:
-            self.database._pool.put(self._session)
+        self.database._pool.put(self._session)
         self._session = None
 
-    def retry_transaction(self):
-        """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:
-            self._spanner_transaction_started = False
-            attempt += 1
-            if attempt > MAX_INTERNAL_RETRIES:
-                raise
-
-            try:
-                self._rerun_previous_statements()
-                break
-            except Aborted as exc:
-                delay = _get_retry_delay(exc.errors[0], attempt)
-                if delay:
-                    time.sleep(delay)
-
-    def _rerun_previous_statements(self):
-        """
-        Helper to run all the remembered statements
-        from the last transaction.
-        """
-        for statement in self._statements:
-            if isinstance(statement, list):
-                statements, checksum = statement
-
-                transaction = self.transaction_checkout()
-                statements_tuple = []
-                for single_statement in statements:
-                    statements_tuple.append(single_statement.get_tuple())
-                status, res = transaction.batch_update(statements_tuple)
-
-                if status.code == ABORTED:
-                    raise Aborted(status.details)
-
-                retried_checksum = ResultsChecksum()
-                retried_checksum.consume_result(res)
-                retried_checksum.consume_result(status.code)
-
-                _compare_checksums(checksum, retried_checksum)
-            else:
-                res_iter, retried_checksum = self.run_statement(statement, retried=True)
-                # executing all the completed statements
-                if statement != self._statements[-1]:
-                    for res in res_iter:
-                        retried_checksum.consume_result(res)
-
-                    _compare_checksums(statement.checksum, retried_checksum)
-                # executing the failed statement
-                else:
-                    # streaming up to the failed result or
-                    # to the end of the streaming iterator
-                    while len(retried_checksum) < len(statement.checksum):
-                        try:
-                            res = next(iter(res_iter))
-                            retried_checksum.consume_result(res)
-                        except StopIteration:
-                            break
-
-                    _compare_checksums(statement.checksum, retried_checksum)
-
     def transaction_checkout(self):
         """Get a Cloud Spanner transaction.
 
@@ -433,12 +393,10 @@ class Connection:
 
     def commit(self):
         """Commits any pending transaction to the database.
-
         This is a no-op if there is no active client transaction.
         """
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
-
         if not self._client_transaction_started:
             warnings.warn(
                 CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
@@ -450,17 +408,13 @@ class Connection:
             if self._spanner_transaction_started and not self._read_only:
                 self._transaction.commit()
         except Aborted:
-            self.retry_transaction()
+            self._transaction_helper.retry_transaction()
             self.commit()
         finally:
-            self._release_session()
-            self._statements = []
-            self._transaction_begin_marked = False
-            self._spanner_transaction_started = False
+            self._reset_post_commit_or_rollback()
 
     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:
@@ -468,15 +422,17 @@ class Connection:
                 CLIENT_TRANSACTION_NOT_STARTED_WARNING, UserWarning, stacklevel=2
             )
             return
-
         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
+            self._reset_post_commit_or_rollback()
+
+    def _reset_post_commit_or_rollback(self):
+        self._release_session()
+        self._transaction_helper.reset()
+        self._transaction_begin_marked = False
+        self._spanner_transaction_started = False
 
     @check_not_closed
     def cursor(self):
@@ -493,7 +449,7 @@ class Connection:
 
             return self.database.update_ddl(ddl_statements).result()
 
-    def run_statement(self, statement: Statement, retried=False):
+    def run_statement(self, statement: Statement):
         """Run single SQL statement in begun transaction.
 
         This method is never used in autocommit mode. In
@@ -513,17 +469,11 @@ class Connection:
                   checksum of this statement results.
         """
         transaction = self.transaction_checkout()
-        if not retried:
-            self._statements.append(statement)
-
-        return (
-            transaction.execute_sql(
-                statement.sql,
-                statement.params,
-                param_types=statement.param_types,
-                request_options=self.request_options,
-            ),
-            ResultsChecksum() if retried else statement.checksum,
+        return transaction.execute_sql(
+            statement.sql,
+            statement.params,
+            param_types=statement.param_types,
+            request_options=self.request_options,
         )
 
     @check_not_closed
@@ -600,15 +550,7 @@ class Connection:
     ):
         statement = parsed_statement.statement
         partitioned_query = parsed_statement.client_side_statement_params[0]
-        if _get_statement_type(Statement(partitioned_query)) is not StatementType.QUERY:
-            raise ProgrammingError(
-                "Only queries can be partitioned. Invalid statement: " + statement.sql
-            )
-        if self.read_only is not True and self._client_transaction_started is True:
-            raise ProgrammingError(
-                "Partitioned query not supported as the connection is not in "
-                "read only mode or ReadWrite transaction started"
-            )
+        self._partitioned_query_validation(partitioned_query, statement)
 
         batch_snapshot = self._database.batch_snapshot()
         partition_ids = []
@@ -620,17 +562,18 @@ class Connection:
                 query_options=query_options,
             )
         )
+
+        batch_transaction_id = batch_snapshot.get_batch_transaction_id()
         for partition in partitions:
-            batch_transaction_id = batch_snapshot.get_batch_transaction_id()
             partition_ids.append(
                 partition_helper.encode_to_string(batch_transaction_id, partition)
             )
         return partition_ids
 
     @check_not_closed
-    def run_partition(self, batch_transaction_id):
+    def run_partition(self, encoded_partition_id):
         partition_id: PartitionId = partition_helper.decode_from_string(
-            batch_transaction_id
+            encoded_partition_id
         )
         batch_transaction_id = partition_id.batch_transaction_id
         batch_snapshot = self._database.batch_snapshot(
@@ -640,6 +583,60 @@ class Connection:
         )
         return batch_snapshot.process(partition_id.partition_result)
 
+    @check_not_closed
+    def run_partitioned_query(
+        self,
+        parsed_statement: ParsedStatement,
+    ):
+        statement = parsed_statement.statement
+        partitioned_query = parsed_statement.client_side_statement_params[0]
+        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
+        )
+
+    @check_not_closed
+    def _set_autocommit_dml_mode(
+        self,
+        parsed_statement: ParsedStatement,
+    ):
+        autocommit_dml_mode_str = parsed_statement.client_side_statement_params[0]
+        autocommit_dml_mode = AutocommitDmlMode[autocommit_dml_mode_str.upper()]
+        self.set_autocommit_dml_mode(autocommit_dml_mode)
+
+    def set_autocommit_dml_mode(
+        self,
+        autocommit_dml_mode,
+    ):
+        """
+        Sets the mode for executing DML statements in autocommit mode for this connection.
+        This mode is only used when the connection is in autocommit mode, and may only
+        be set while the transaction is in autocommit mode and not in a temporary transaction.
+        """
+
+        if self._client_transaction_started is True:
+            raise ProgrammingError(
+                "Cannot set autocommit DML mode while not in autocommit mode or while a transaction is active."
+            )
+        if self.read_only is True:
+            raise ProgrammingError(
+                "Cannot set autocommit DML mode for a read-only connection."
+            )
+        if self._batch_mode is not BatchMode.NONE:
+            raise ProgrammingError("Cannot set autocommit DML mode while in a batch.")
+        self._autocommit_dml_mode = autocommit_dml_mode
+
+    def _partitioned_query_validation(self, partitioned_query, statement):
+        if _get_statement_type(Statement(partitioned_query)) is not StatementType.QUERY:
+            raise ProgrammingError(
+                "Only queries can be partitioned. Invalid statement: " + statement.sql
+            )
+        if self.read_only is not True and self._client_transaction_started is True:
+            raise ProgrammingError(
+                "Partitioned query is not supported, because the connection is in a read/write transaction."
+            )
+
     def __enter__(self):
         return self
 
@@ -729,9 +726,10 @@ def connect(
             raise ValueError("project in url does not match client object project")
 
     instance = client.instance(instance_id)
-    conn = Connection(
-        instance, instance.database(database_id, pool=pool) if database_id else None
-    )
+    database = None
+    if database_id:
+        database = instance.database(database_id, pool=pool)
+    conn = Connection(instance, database)
     if pool is not None:
         conn._own_pool = False