Automatically set SQL driver if unset.

tiled/_tests/conftest.py

@@ -13,6 +13,7 @@
 from ..catalog import from_uri, in_memory
 from ..client.base import BaseClient
 from ..server.settings import get_settings
+from ..utils import ensure_specified_sql_driver
 from .utils import enter_password as utils_enter_password
 from .utils import temp_postgres
 
@@ -152,7 +153,7 @@ async def postgresql_with_example_data_adapter(request, tmpdir):
     if uri.endswith("/"):
         uri = uri[:-1]
     uri_with_database_name = f"{uri}/{DATABASE_NAME}"
-    engine = create_async_engine(uri_with_database_name)
+    engine = create_async_engine(ensure_specified_sql_driver(uri_with_database_name))
     try:
         async with engine.connect():
             pass

tiled/_tests/utils.py

@@ -14,6 +14,7 @@
 
 from ..client import context
 from ..client.base import BaseClient
+from ..utils import ensure_specified_sql_driver
 
 if sys.version_info < (3, 9):
     import importlib_resources as resources
@@ -33,7 +34,7 @@ async def temp_postgres(uri):
     if uri.endswith("/"):
         uri = uri[:-1]
     # Create a fresh database.
-    engine = create_async_engine(uri)
+    engine = create_async_engine(ensure_specified_sql_driver(uri))
     database_name = f"tiled_test_disposable_{uuid.uuid4().hex}"
     async with engine.connect() as connection:
         await connection.execute(

tiled/authn_database/connection_pool.py

@@ -2,6 +2,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 
 from ..server.settings import get_settings
+from ..utils import ensure_specified_sql_driver
 
 # A given process probably only has one of these at a time, but we
 # key on database_settings just case in some testing context or something
@@ -16,7 +17,9 @@ def open_database_connection_pool(database_settings):
     # kwargs["pool_pre_ping"] = database_settings.pool_pre_ping
     # kwargs["max_overflow"] = database_settings.max_overflow
     engine = create_async_engine(
-        database_settings.uri, connect_args=connect_args, **kwargs
+        ensure_specified_sql_driver(database_settings.uri),
+        connect_args=connect_args,
+        **kwargs,
     )
     _connection_pools[database_settings] = engine
     return engine

tiled/catalog/adapter.py

@@ -70,6 +70,7 @@
     OneShotCachedMap,
     UnsupportedQueryType,
     ensure_awaitable,
+    ensure_specified_sql_driver,
     ensure_uri,
     import_object,
     path_from_uri,
@@ -1381,7 +1382,10 @@ def from_uri(
     else:
         poolclass = None  # defer to sqlalchemy default
     engine = create_async_engine(
-        uri, echo=echo, json_serializer=json_serializer, poolclass=poolclass
+        ensure_specified_sql_driver(uri),
+        echo=echo,
+        json_serializer=json_serializer,
+        poolclass=poolclass,
     )
     if engine.dialect.name == "sqlite":
         event.listens_for(engine.sync_engine, "connect")(_set_sqlite_pragma)

tiled/commandline/_admin.py

@@ -27,9 +27,10 @@ def initialize_database(database_uri: str):
         REQUIRED_REVISION,
         initialize_database,
     )
+    from ..utils import ensure_specified_sql_driver
 
     async def do_setup():
-        engine = create_async_engine(database_uri)
+        engine = create_async_engine(ensure_specified_sql_driver(database_uri))
         redacted_url = engine.url._replace(password="[redacted]")
         try:
             await check_database(engine, REQUIRED_REVISION, ALL_REVISIONS)
@@ -71,9 +72,10 @@ def upgrade_database(
         ALEMBIC_INI_TEMPLATE_PATH,
     )
     from ..authn_database.core import ALL_REVISIONS
+    from ..utils import ensure_specified_sql_driver
 
     async def do_setup():
-        engine = create_async_engine(database_uri)
+        engine = create_async_engine(ensure_specified_sql_driver(database_uri))
         redacted_url = engine.url._replace(password="[redacted]")
         current_revision = await get_current_revision(engine, ALL_REVISIONS)
         await engine.dispose()
@@ -107,9 +109,10 @@ def downgrade_database(
         ALEMBIC_INI_TEMPLATE_PATH,
     )
     from ..authn_database.core import ALL_REVISIONS
+    from ..utils import ensure_specified_sql_driver
 
     async def do_setup():
-        engine = create_async_engine(database_uri)
+        engine = create_async_engine(ensure_specified_sql_driver(database_uri))
         redacted_url = engine.url._replace(password="[redacted]")
         current_revision = await get_current_revision(engine, ALL_REVISIONS)
         if current_revision is None:

tiled/commandline/_catalog.py

@@ -44,14 +44,14 @@ def init(
     from ..alembic_utils import UninitializedDatabase, check_database, stamp_head
     from ..catalog.alembic_constants import ALEMBIC_DIR, ALEMBIC_INI_TEMPLATE_PATH
     from ..catalog.core import ALL_REVISIONS, REQUIRED_REVISION, initialize_database
-    from ..utils import SCHEME_PATTERN
+    from ..utils import SCHEME_PATTERN, ensure_specified_sql_driver
 
     if not SCHEME_PATTERN.match(database):
         # Interpret URI as filepath.
         database = f"sqlite+aiosqlite:///{database}"
 
     async def do_setup():
-        engine = create_async_engine(database)
+        engine = create_async_engine(ensure_specified_sql_driver(database))
         redacted_url = engine.url._replace(password="[redacted]")
         try:
             await check_database(engine, REQUIRED_REVISION, ALL_REVISIONS)
@@ -94,9 +94,10 @@ def upgrade_database(
     from ..alembic_utils import get_current_revision, upgrade
     from ..catalog.alembic_constants import ALEMBIC_DIR, ALEMBIC_INI_TEMPLATE_PATH
     from ..catalog.core import ALL_REVISIONS
+    from ..utils import ensure_specified_sql_driver
 
     async def do_setup():
-        engine = create_async_engine(database_uri)
+        engine = create_async_engine(ensure_specified_sql_driver(database_uri))
         redacted_url = engine.url._replace(password="[redacted]")
         current_revision = await get_current_revision(engine, ALL_REVISIONS)
         await engine.dispose()
@@ -127,9 +128,10 @@ def downgrade_database(
     from ..alembic_utils import downgrade, get_current_revision
     from ..catalog.alembic_constants import ALEMBIC_DIR, ALEMBIC_INI_TEMPLATE_PATH
     from ..catalog.core import ALL_REVISIONS
+    from ..utils import ensure_specified_sql_driver
 
     async def do_setup():
-        engine = create_async_engine(database_uri)
+        engine = create_async_engine(ensure_specified_sql_driver(database_uri))
         redacted_url = engine.url._replace(password="[redacted]")
         current_revision = await get_current_revision(engine, ALL_REVISIONS)
         if current_revision is None:

tiled/commandline/_serve.py

@@ -129,8 +129,9 @@ def serve_directory(
     from ..alembic_utils import stamp_head
     from ..catalog.alembic_constants import ALEMBIC_DIR, ALEMBIC_INI_TEMPLATE_PATH
     from ..catalog.core import initialize_database
+    from ..utils import ensure_specified_sql_driver
 
-    engine = create_async_engine(database)
+    engine = create_async_engine(ensure_specified_sql_driver(database))
     asyncio.run(initialize_database(engine))
     stamp_head(ALEMBIC_INI_TEMPLATE_PATH, ALEMBIC_DIR, database)
 
@@ -389,8 +390,9 @@ def serve_catalog(
         from ..alembic_utils import stamp_head
         from ..catalog.alembic_constants import ALEMBIC_DIR, ALEMBIC_INI_TEMPLATE_PATH
         from ..catalog.core import initialize_database
+        from ..utils import ensure_specified_sql_driver
 
-        engine = create_async_engine(database)
+        engine = create_async_engine(ensure_specified_sql_driver(database))
         asyncio.run(initialize_database(engine))
         stamp_head(ALEMBIC_INI_TEMPLATE_PATH, ALEMBIC_DIR, database)
 

tiled/utils.py

@@ -721,6 +721,31 @@ def ensure_uri(uri_or_path) -> str:
     return str(uri_str)
 
 
+SCHEME_TO_SCHEME_PLUS_DRIVER = {
+    "postgresql": "postgresql+asyncpg",
+    "sqlite": "sqlite+aiosqlite",
+}
+
+
+def ensure_specified_sql_driver(uri: str) -> str:
+    """
+    Given a URI without a driver in the scheme, add Tiled's preferred driver.
+
+    If a driver is already specified, the specified one will be used; it
+    will NOT be overriden by this function.
+
+    'postgresql://...' -> 'postgresql+asynpg://...'
+    'sqlite://...' -> 'sqlite+aiosqlite://...'
+    'postgresql+asyncpg://...' -> 'postgresql+asynpg://...'
+    'postgresql+my_custom_driver://...' -> 'postgresql+my_custom_driver://...'
+    """
+    parsed_uri = urlparse(uri)
+    scheme = parsed_uri.scheme
+    new_scheme = SCHEME_TO_SCHEME_PLUS_DRIVER.get(scheme, scheme)
+    updated_uri = urlunparse(parsed_uri._replace(scheme=new_scheme))
+    return updated_uri
+
+
 class catch_warning_msg(warnings.catch_warnings):
     """Backward compatible version of catch_warnings for python <3.11.