Merge branch 'main' into support-db-roles-in-connect

Author: olavloite

google/cloud/spanner_dbapi/connection.py

@@ -805,7 +805,7 @@ def connect(
         database = instance.database(
             database_id, pool=pool, database_role=database_role
         )
-    conn = Connection(instance, database)
+    conn = Connection(instance, database, **kwargs)
     if pool is not None:
         conn._own_pool = False
 

tests/system/test_dbapi.py

@@ -865,9 +865,9 @@ def test_execute_batch_dml_abort_retry(self, dbapi_database):
         self._cursor.execute("run batch")
         dbapi_database._method_abort_interceptor.reset()
         self._conn.commit()
-        assert method_count_interceptor._counts[COMMIT_METHOD] == 1
-        assert method_count_interceptor._counts[EXECUTE_BATCH_DML_METHOD] == 3
-        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] == 6
+        assert method_count_interceptor._counts[COMMIT_METHOD] >= 1
+        assert method_count_interceptor._counts[EXECUTE_BATCH_DML_METHOD] >= 3
+        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] >= 6
 
         self._cursor.execute("SELECT * FROM contacts")
         got_rows = self._cursor.fetchall()
@@ -879,28 +879,28 @@ def test_multiple_aborts_in_transaction(self, dbapi_database):
 
         method_count_interceptor = dbapi_database._method_count_interceptor
         method_count_interceptor.reset()
-        # called 3 times
+        # called at least 3 times
         self._insert_row(1)
         dbapi_database._method_abort_interceptor.set_method_to_abort(
             EXECUTE_STREAMING_SQL_METHOD, self._conn
         )
-        # called 3 times
+        # called at least 3 times
         self._cursor.execute("SELECT * FROM contacts")
         dbapi_database._method_abort_interceptor.reset()
         self._cursor.fetchall()
-        # called 2 times
+        # called at least 2 times
         self._insert_row(2)
-        # called 2 times
+        # called at least 2 times
         self._cursor.execute("SELECT * FROM contacts")
         self._cursor.fetchone()
         dbapi_database._method_abort_interceptor.set_method_to_abort(
             COMMIT_METHOD, self._conn
         )
-        # called 2 times
+        # called at least 2 times
         self._conn.commit()
         dbapi_database._method_abort_interceptor.reset()
-        assert method_count_interceptor._counts[COMMIT_METHOD] == 2
-        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] == 10
+        assert method_count_interceptor._counts[COMMIT_METHOD] >= 2
+        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] >= 10
 
         self._cursor.execute("SELECT * FROM contacts")
         got_rows = self._cursor.fetchall()
@@ -921,8 +921,8 @@ def test_consecutive_aborted_transactions(self, dbapi_database):
         )
         self._conn.commit()
         dbapi_database._method_abort_interceptor.reset()
-        assert method_count_interceptor._counts[COMMIT_METHOD] == 2
-        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] == 6
+        assert method_count_interceptor._counts[COMMIT_METHOD] >= 2
+        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] >= 6
 
         method_count_interceptor = dbapi_database._method_count_interceptor
         method_count_interceptor.reset()
@@ -935,8 +935,8 @@ def test_consecutive_aborted_transactions(self, dbapi_database):
         )
         self._conn.commit()
         dbapi_database._method_abort_interceptor.reset()
-        assert method_count_interceptor._counts[COMMIT_METHOD] == 2
-        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] == 6
+        assert method_count_interceptor._counts[COMMIT_METHOD] >= 2
+        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] >= 6
 
         self._cursor.execute("SELECT * FROM contacts")
         got_rows = self._cursor.fetchall()

tests/unit/spanner_dbapi/test_connect.py

@@ -137,3 +137,17 @@ def test_w_credential_file_path(self, mock_client):
         client_info = factory.call_args_list[0][1]["client_info"]
         self.assertEqual(client_info.user_agent, USER_AGENT)
         self.assertEqual(client_info.python_version, PY_VERSION)
+
+    def test_with_kwargs(self, mock_client):
+        from google.cloud.spanner_dbapi import connect
+        from google.cloud.spanner_dbapi import Connection
+
+        client = mock_client.return_value
+        instance = client.instance.return_value
+        database = instance.database.return_value
+        self.assertIsNotNone(database)
+
+        connection = connect(INSTANCE, DATABASE, ignore_transaction_warnings=True)
+
+        self.assertIsInstance(connection, Connection)
+        self.assertTrue(connection._ignore_transaction_warnings)