Merge branch 'main' into test-json-array

Author: olavloite

.kokoro/requirements.txt

@@ -4,9 +4,9 @@
 #
 #    pip-compile --generate-hashes requirements.in
 #
-argcomplete==3.5.1 \
-    --hash=sha256:1a1d148bdaa3e3b93454900163403df41448a248af01b6e849edc5ac08e6c363 \
-    --hash=sha256:eb1ee355aa2557bd3d0145de7b06b2a45b0ce461e1e7813f5d066039ab4177b4
+argcomplete==3.5.2 \
+    --hash=sha256:036d020d79048a5d525bc63880d7a4b8d1668566b8a76daf1144c0bbe0f63472 \
+    --hash=sha256:23146ed7ac4403b70bd6026402468942ceba34a6732255b9edf5b7354f68a6bb
     # via nox
 attrs==24.2.0 \
     --hash=sha256:5cfb1b9148b5b086569baec03f20d7b6bf3bcacc9a42bebf87ffaaca362f6346 \

requirements.txt

@@ -250,9 +250,9 @@ importlib-metadata==8.5.0 \
     --hash=sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b \
     --hash=sha256:71522656f0abace1d072b9e5481a48f07c138e00f079c38c8f883823f9c26bd7
     # via opentelemetry-api
-Mako==1.3.7 \
-    --hash=sha256:20405b1232e0759f0e7d87b01f6bb94fce0761747f1cb876ecf90bd512d0b639 \
-    --hash=sha256:d18f990ad57f800ce8e76cbfb0b74afe471c293517e9f5003ace6dad5aa72c36
+Mako==1.3.8 \
+    --hash=sha256:42f48953c7eb91332040ff567eb7eea69b22e7a4affbc5ba8e845e8f730f6627 \
+    --hash=sha256:577b97e414580d3e088d47c2dbbe9594aa7a5146ed2875d4dfa9075af2dd3cc8
     # via alembic
 markupsafe==3.0.2 \
     --hash=sha256:0bff5e0ae4ef2e1ae4fdf2dfd5b76c75e5c2fa4132d05fc1b0dabcd20c7e28c4 \

samples/partitioned_dml_sample.py

@@ -0,0 +1,45 @@
+# Copyright 2024 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 google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
+from sqlalchemy import create_engine, text
+
+from sample_helper import run_sample
+
+# Shows how to use Partitioned DML using SQLAlchemy and Spanner.
+def partitioned_dml_sample():
+    engine = create_engine(
+        "spanner:///projects/sample-project/"
+        "instances/sample-instance/"
+        "databases/sample-database",
+        echo=True,
+    )
+    # Get a connection in auto-commit mode.
+    # Partitioned DML can only be executed in auto-commit mode, as each
+    # Partitioned DML transaction can only consist of one statement.
+    with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as connection:
+        # Set the DML mode to PARTITIONED_NON_ATOMIC.
+        connection.connection.set_autocommit_dml_mode(
+            AutocommitDmlMode.PARTITIONED_NON_ATOMIC
+        )
+        # Use a bulk update statement to back-fill a column.
+        lower_bound_rowcount = connection.execute(
+            text("update venues set active=true where active is null")
+        ).rowcount
+        # Partitioned DML returns the lower-bound update count.
+        print("Updated at least ", lower_bound_rowcount, " venue records")
+
+
+if __name__ == "__main__":
+    run_sample(partitioned_dml_sample)

test/mockserver_tests/mock_server_test_base.py

@@ -11,6 +11,7 @@
 # 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 google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
 from sqlalchemy import Engine, create_engine
 from sqlalchemy.testing.plugin.plugin_base import fixtures
 import google.cloud.spanner_v1.types.type as spanner_type
@@ -35,16 +36,14 @@ def add_result(sql: str, result: ResultSet):
     MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result)
 
 
-def add_update_count(sql: str, count: int):
-    result = result_set.ResultSet(
-        dict(
-            stats=result_set.ResultSetStats(
-                dict(
-                    row_count_exact=count,
-                )
-            ),
-        )
-    )
+def add_update_count(
+    sql: str, count: int, dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
+):
+    if dml_mode == AutocommitDmlMode.PARTITIONED_NON_ATOMIC:
+        stats = dict(row_count_lower_bound=count)
+    else:
+        stats = dict(row_count_exact=count)
+    result = result_set.ResultSet(dict(stats=result_set.ResultSetStats(stats)))
     add_result(sql, result)
 
 

test/mockserver_tests/mock_spanner.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from google.cloud.spanner_v1 import TransactionOptions, ResultSetMetadata
+from google.cloud.spanner_v1 import (
+    TransactionOptions,
+    ResultSetMetadata,
+    ExecuteSqlRequest,
+)
 from google.protobuf import empty_pb2
 import test.mockserver_tests.spanner_pb2_grpc as spanner_grpc
 import test.mockserver_tests.spanner_database_admin_pb2_grpc as database_admin_grpc
@@ -40,23 +44,25 @@ def get_result(self, sql: str) -> result_set.ResultSet:
         return result
 
     def get_result_as_partial_result_sets(
-        self, sql: str
+        self, sql: str, started_transaction: transaction.Transaction
     ) -> [result_set.PartialResultSet]:
         result: result_set.ResultSet = self.get_result(sql)
         partials = []
         first = True
         if len(result.rows) == 0:
             partial = result_set.PartialResultSet()
-            partial.metadata = result.metadata
+            partial.metadata = ResultSetMetadata(result.metadata)
             partials.append(partial)
         else:
             for row in result.rows:
                 partial = result_set.PartialResultSet()
                 if first:
-                    partial.metadata = result.metadata
+                    partial.metadata = ResultSetMetadata(result.metadata)
                 partial.values.extend(row)
                 partials.append(partial)
         partials[len(partials) - 1].stats = result.stats
+        if started_transaction:
+            partials[0].metadata.transaction = started_transaction
         return partials
 
 
@@ -120,9 +126,16 @@ def ExecuteSql(self, request, context):
         self._requests.append(request)
         return result_set.ResultSet()
 
-    def ExecuteStreamingSql(self, request, context):
+    def ExecuteStreamingSql(self, request: ExecuteSqlRequest, context):
         self._requests.append(request)
-        partials = self.mock_spanner.get_result_as_partial_result_sets(request.sql)
+        started_transaction = None
+        if not request.transaction.begin == TransactionOptions():
+            started_transaction = self.__create_transaction(
+                request.session, request.transaction.begin
+            )
+        partials = self.mock_spanner.get_result_as_partial_result_sets(
+            request.sql, started_transaction
+        )
         for result in partials:
             yield result
 

test/mockserver_tests/test_basics.py

@@ -13,7 +13,17 @@
 # limitations under the License.
 
 from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
-from sqlalchemy import create_engine, select, MetaData, Table, Column, Integer, String
+from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
+from sqlalchemy import (
+    create_engine,
+    select,
+    MetaData,
+    Table,
+    Column,
+    Integer,
+    String,
+    text,
+)
 from sqlalchemy.testing import eq_, is_instance_of
 from google.cloud.spanner_v1 import (
     FixedSizePool,
@@ -26,6 +36,7 @@
     MockServerTestBase,
     add_select1_result,
     add_result,
+    add_update_count,
 )
 
 
@@ -127,3 +138,21 @@ def test_create_multiple_tables(self):
                 "\n) PRIMARY KEY (id)",
                 requests[0].statements[i],
             )
+
+    def test_partitioned_dml(self):
+        sql = "UPDATE singers SET checked=true WHERE active = true"
+        add_update_count(sql, 100, AutocommitDmlMode.PARTITIONED_NON_ATOMIC)
+        engine = create_engine(
+            "spanner:///projects/p/instances/i/databases/d",
+            connect_args={"client": self.client, "pool": PingingPool(size=10)},
+        )
+        # TODO: Support autocommit_dml_mode as a connection variable in execution
+        #       options.
+        with engine.connect().execution_options(
+            isolation_level="AUTOCOMMIT"
+        ) as connection:
+            connection.connection.set_autocommit_dml_mode(
+                AutocommitDmlMode.PARTITIONED_NON_ATOMIC
+            )
+            results = connection.execute(text(sql)).rowcount
+            eq_(100, results)