chore: update tests for mux session usage

Author: olavloite

.gitignore

@@ -48,6 +48,7 @@ docs/_build
 docs.metadata
 
 # Virtual environment
+.venv/
 env/
 coverage.xml
 sponge_log.xml

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import base64
+import logging
 
 import re
 
@@ -927,6 +928,18 @@ def _get_table_filter_query(
 
         return table_filter_query
 
+    def connect(self, *cargs, **cparams):
+        conn = super().connect(*cargs, **cparams)
+        database = getattr(conn, "_database")
+        if database:
+            logger = cparams.get("logger")
+            if logger is None:
+                logger = logging.getLogger()
+                database._logger = logger
+            else:
+                database._logger = logger
+        return conn
+
     def create_connect_args(self, url):
         """Parse connection args from the given URL.
 

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.
+import logging
 
 from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
 from sqlalchemy import Engine, create_engine
@@ -22,7 +23,6 @@
 from google.cloud.spanner_v1 import (
     Client,
     ResultSet,
-    PingingPool,
     TypeCode,
 )
 from google.cloud.spanner_v1.database import Database
@@ -151,6 +151,7 @@ def setup_method(self):
         self._client = None
         self._instance = None
         self._database = None
+        _ = self.database
 
     def teardown_method(self):
         MockServerTestBase.spanner_service.clear_requests()
@@ -159,7 +160,7 @@ def teardown_method(self):
     def create_engine(self) -> Engine:
         return create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": PingingPool(size=10)},
+            connect_args={"client": self.client},
         )
 
     @property
@@ -177,13 +178,13 @@ def client(self) -> Client:
     @property
     def instance(self) -> Instance:
         if self._instance is None:
-            self._instance = self.client.instance("test-instance")
+            self._instance = self.client.instance("i")
         return self._instance
 
     @property
     def database(self) -> Database:
+        logger = logging.getLogger("level warning")
+        logger.setLevel(logging.WARN)
         if self._database is None:
-            self._database = self.instance.database(
-                "test-database", pool=PingingPool(size=10)
-            )
+            self._database = self.instance.database("d", logger=logger)
         return self._database

test/mockserver_tests/test_auto_increment.py

@@ -16,9 +16,8 @@
 from sqlalchemy.orm import Session
 from sqlalchemy.testing import eq_, is_instance_of
 from google.cloud.spanner_v1 import (
-    FixedSizePool,
     ResultSet,
-    BatchCreateSessionsRequest,
+    CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
     BeginTransactionRequest,
@@ -47,7 +46,7 @@ def test_create_table(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         Base.metadata.create_all(engine)
         requests = self.database_admin_service.requests
@@ -76,7 +75,7 @@ def test_create_auto_increment_table(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         engine.dialect.use_auto_increment = True
         Base.metadata.create_all(engine)
@@ -105,7 +104,7 @@ def test_create_table_with_specific_sequence_kind(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         engine.dialect.default_sequence_kind = "non_existing_kind"
         Base.metadata.create_all(engine)
@@ -128,7 +127,7 @@ def test_insert_row(self):
         self.add_insert_result("INSERT INTO singers (name) VALUES (@a0) THEN RETURN id")
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
 
         with Session(engine) as session:
@@ -141,7 +140,7 @@ def test_insert_row(self):
         # Verify the requests that we got.
         requests = self.spanner_service.requests
         eq_(4, len(requests))
-        is_instance_of(requests[0], BatchCreateSessionsRequest)
+        is_instance_of(requests[0], CreateSessionRequest)
         is_instance_of(requests[1], BeginTransactionRequest)
         is_instance_of(requests[2], ExecuteSqlRequest)
         is_instance_of(requests[3], CommitRequest)
@@ -154,7 +153,7 @@ def test_insert_row_with_pk_value(self):
         add_update_count("INSERT INTO singers (id, name) VALUES (@a0, @a1)", 1)
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
 
         with Session(engine) as session:

test/mockserver_tests/test_basics.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import datetime
+
 from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
 from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
 from sqlalchemy import (
@@ -32,8 +33,7 @@
 from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.testing import eq_, is_instance_of
 from google.cloud.spanner_v1 import (
-    FixedSizePool,
-    BatchCreateSessionsRequest,
+    CreateSessionRequest,
     ExecuteSqlRequest,
     ResultSet,
     PingingPool,
@@ -58,7 +58,7 @@ def verify_select1(self, results):
         eq_(1, len(result_list))
         requests = self.spanner_service.requests
         eq_(2, len(requests))
-        is_instance_of(requests[0], BatchCreateSessionsRequest)
+        is_instance_of(requests[0], CreateSessionRequest)
         is_instance_of(requests[1], ExecuteSqlRequest)
 
     def test_select1(self):
@@ -71,7 +71,7 @@ def test_sqlalchemy_select1(self):
         add_select1_result()
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": PingingPool(size=10)},
+            connect_args={"client": self.client},
         )
         with engine.connect().execution_options(
             isolation_level="AUTOCOMMIT"
@@ -113,7 +113,7 @@ def test_create_table(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         metadata = MetaData()
         Table(
@@ -150,7 +150,7 @@ def test_create_table_in_schema(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         metadata = MetaData()
         Table(
@@ -192,7 +192,7 @@ def test_create_multiple_tables(self):
             )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         metadata = MetaData()
         for i in range(2):
@@ -260,7 +260,7 @@ class Singer(Base):
 
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
 
         with Session(engine) as session:
@@ -277,18 +277,18 @@ def test_database_role(self):
             "spanner:///projects/p/instances/i/databases/d",
             connect_args={
                 "client": self.client,
-                "pool": FixedSizePool(size=10),
+                # "pool": FixedSizePool(size=10),
                 "database_role": "my_role",
             },
         )
         with Session(engine.execution_options(isolation_level="autocommit")) as session:
             session.execute(select(1))
         requests = self.spanner_service.requests
         eq_(2, len(requests))
-        is_instance_of(requests[0], BatchCreateSessionsRequest)
+        is_instance_of(requests[0], CreateSessionRequest)
         is_instance_of(requests[1], ExecuteSqlRequest)
-        request: BatchCreateSessionsRequest = requests[0]
-        eq_("my_role", request.session_template.creator_role)
+        request: CreateSessionRequest = requests[0]
+        eq_("my_role", request.session.creator_role)
 
     def test_select_table_in_named_schema(self):
         class Base(DeclarativeBase):
@@ -311,7 +311,7 @@ class Singer(Base):
         add_singer_query_result(query)
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
 
         insert = "INSERT INTO my_schema.singers (name) VALUES (@a0) THEN RETURN id"

test/mockserver_tests/test_bit_reversed_sequence.py

@@ -16,9 +16,8 @@
 from sqlalchemy.orm import Session
 from sqlalchemy.testing import eq_, is_instance_of
 from google.cloud.spanner_v1 import (
-    FixedSizePool,
     ResultSet,
-    BatchCreateSessionsRequest,
+    CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
     BeginTransactionRequest,
@@ -54,7 +53,7 @@ def test_create_table(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         Base.metadata.create_all(engine)
         requests = self.database_admin_service.requests
@@ -115,7 +114,7 @@ def test_insert_row(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
 
         with Session(engine) as session:
@@ -128,7 +127,7 @@ def test_insert_row(self):
         # Verify the requests that we got.
         requests = self.spanner_service.requests
         eq_(4, len(requests))
-        is_instance_of(requests[0], BatchCreateSessionsRequest)
+        is_instance_of(requests[0], CreateSessionRequest)
         is_instance_of(requests[1], BeginTransactionRequest)
         is_instance_of(requests[2], ExecuteSqlRequest)
         is_instance_of(requests[3], CommitRequest)

test/mockserver_tests/test_commit_timestamp.py

@@ -14,10 +14,7 @@
 
 from sqlalchemy import create_engine
 from sqlalchemy.testing import eq_, is_instance_of
-from google.cloud.spanner_v1 import (
-    FixedSizePool,
-    ResultSet,
-)
+from google.cloud.spanner_v1 import ResultSet
 from test.mockserver_tests.mock_server_test_base import (
     MockServerTestBase,
     add_result,
@@ -47,7 +44,7 @@ def test_create_table(self):
         )
         engine = create_engine(
             "spanner:///projects/p/instances/i/databases/d",
-            connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
+            connect_args={"client": self.client},
         )
         Base.metadata.create_all(engine)
         requests = self.database_admin_service.requests

test/mockserver_tests/test_float32.py

@@ -19,7 +19,7 @@
     is_false,
 )
 from google.cloud.spanner_v1 import (
-    BatchCreateSessionsRequest,
+    CreateSessionRequest,
     ExecuteSqlRequest,
     ResultSet,
     ResultSetStats,
@@ -59,7 +59,7 @@ def test_insert_data(self):
 
             requests = self.spanner_service.requests
             eq_(4, len(requests))
-            is_instance_of(requests[0], BatchCreateSessionsRequest)
+            is_instance_of(requests[0], CreateSessionRequest)
             is_instance_of(requests[1], BeginTransactionRequest)
             is_instance_of(requests[2], ExecuteSqlRequest)
             is_instance_of(requests[3], CommitRequest)