Fix SQLAlchemy tests

Author: CarlosPaparoni

baseplate/clients/sqlalchemy.py

@@ -10,7 +10,7 @@
 from sqlalchemy import create_engine, event
 from sqlalchemy.engine import Connection, Engine, ExceptionContext
 from sqlalchemy.engine.interfaces import ExecutionContext
-from sqlalchemy.engine.url import make_url
+from sqlalchemy.engine.url import URL, make_url
 from sqlalchemy.orm import Session
 from sqlalchemy.pool import QueuePool
 
@@ -87,13 +87,16 @@ def engine_from_config(
             raise TypeError("'secrets' is required if 'credentials_secret' is set")
         credentials = secrets.get_credentials(options.credentials_secret)
 
-        # support sqlalchemy 1.4+ where URL is immutable
-        # https://docs.sqlalchemy.org/en/14/changelog/migration_14.html#the-url-object-is-now-immutable
-        if hasattr(url, "set"):
-            url = url.set(username=credentials.username, password=credentials.password)
-        else:
-            url.username = credentials.username
-            url.password = credentials.password
+        # Create new URL with SQLAlchemy 2.0 compatible parameters
+        url = URL.create(
+            drivername=url.drivername,
+            username=credentials.username,
+            password=credentials.password,
+            host=url.host,
+            port=url.port,
+            database=url.database,
+            query=url.query or {},
+        )
 
     return create_engine(url, **kwargs)
 

tests/integration/message_queue_tests.py

@@ -42,7 +42,7 @@ def test_get_timeout(self):
             with self.assertRaises(TimedOutError):
                 mq.get(timeout=0.1)
             elapsed = time.time() - start
-            self.assertAlmostEqual(elapsed, 0.1, places=2)
+            self.assertAlmostEqual(elapsed, 0.1, places=1)
 
     def test_put_timeout(self):
         message_queue = MessageQueue(self.qname, max_messages=1, max_message_size=1)

tests/integration/sqlalchemy_tests.py

@@ -1,7 +1,7 @@
 import unittest
 
 try:
-    from sqlalchemy import Column, Integer, String
+    from sqlalchemy import Column, Integer, String, text
     from sqlalchemy.dialects.sqlite import BOOLEAN
     from sqlalchemy.exc import OperationalError, StatementError
     from sqlalchemy.ext.declarative import declarative_base
@@ -46,7 +46,8 @@ def setUp(self):
 
     def test_simple_query(self):
         with self.server_span:
-            self.server_span.context.db.execute("SELECT * FROM test;")
+            with self.server_span.context.db.connect() as conn:
+                conn.execute(text("SELECT * FROM test;"))
 
         server_span_observer = self.baseplate_observer.get_only_child()
         span_observer = server_span_observer.get_only_child()
@@ -57,7 +58,8 @@ def test_simple_query(self):
 
     def test_very_long_query(self):
         with self.server_span:
-            self.server_span.context.db.execute("SELECT *" + (" " * 1024) + "FROM test;")
+            with self.server_span.context.db.connect() as conn:
+                conn.execute(text("SELECT *" + (" " * 1024) + "FROM test;"))
 
         server_span_observer = self.baseplate_observer.get_only_child()
         span_observer = server_span_observer.get_only_child()
@@ -72,7 +74,8 @@ def test_nested_local_span(self):
             with self.server_span.make_child(
                 "local", local=True, component_name="example"
             ) as local:
-                local.context.db.execute("SELECT * FROM test;")
+                with local.context.db.connect() as conn:
+                    conn.execute(text("SELECT * FROM test;"))
 
         server_span_observer = self.baseplate_observer.get_only_child()
         local_span_observer = server_span_observer.get_only_child()
@@ -85,7 +88,8 @@ def test_nested_local_span(self):
     def test_error_in_query(self):
         with self.server_span:
             with self.assertRaises(OperationalError):
-                self.context.db.execute("SELECT * FROM does_not_exist;")
+                with self.context.db.connect() as conn:
+                    conn.execute(text("SELECT * FROM does_not_exist;"))
 
         server_span_observer = self.baseplate_observer.get_only_child()
         span_observer = server_span_observer.get_only_child()
@@ -138,4 +142,4 @@ def test_simple_config(self):
 
         context = baseplate.make_context_object()
         with baseplate.make_server_span(context, "test"):
-            context.db.execute("SELECT 1;")
+            context.db.execute(text("SELECT 1;"))

tests/unit/clients/sqlalchemy_tests.py

@@ -30,7 +30,18 @@ def setUp(self):
 
     def test_url(self):
         engine = engine_from_config({"database.url": "sqlite://"})
-        self.assertEqual(engine.url, URL("sqlite"))
+        self.assertEqual(
+            engine.url,
+            URL.create(
+                "sqlite",
+                username=None,
+                password=None,
+                host=None,
+                port=None,
+                database=None,
+                query={},
+            ),
+        )
 
     @mock.patch("baseplate.clients.sqlalchemy.create_engine")
     def test_credentials(self, create_engine_mock):
@@ -45,13 +56,14 @@ def test_credentials(self, create_engine_mock):
             self.secrets,
         )
         create_engine_mock.assert_called_once_with(
-            URL(
+            URL.create(
                 drivername="postgresql",
                 username="reddit",
                 password="password",
                 host="localhost",
                 port="9000",
                 database="db",
+                query={},
             ),
             pool_recycle=60,
             pool_size=10,