update ruff, cleanup database methods on testclient (#63)

* update ruff, cleanup database_exist method on testclient

Changes:

- ruff is now version 0.10.0
- ruff versions of precommit and hatch fmt use the same rules
- apply ruff rules
- cleanup in DatabaseTestClient the database_exist method
- the database_exist method is now more lenient against missing
  permission errors or missing databases for all dbs

* remove dependency on sqlalchemy_utils, improve testclient methods

* fix typings, readd sqlalchemy_utils as debug dependency

* use mysql special user root when possible

* add quoting tests

* harden against segfaults

Author: devkral

.github/workflows/test-suite.yml

@@ -93,4 +93,5 @@ jobs:
         run: "hatch run test:check_types"
       - name: "Run tests"
         if: steps.filters.outputs.src == 'true' || steps.filters.outputs.workflows == 'true' || github.event.schedule != ''
-        run: env TEST_NO_RISK_SEGFAULTS=true hatch test
+        # test with xdist as still segfaults occur
+        run: env TEST_NO_RISK_SEGFAULTS=true hatch test -- -n 1 --dist no

.pre-commit-config.yaml

@@ -13,7 +13,7 @@ repos:
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.6.4
+    rev: v0.10.0
     hooks:
       - id: ruff
         args: ["--fix"]

databasez/__init__.py

@@ -1,5 +1,5 @@
 from databasez.core import Database, DatabaseURL
 
-__version__ = "0.11.1"
+__version__ = "0.11.2"
 
 __all__ = ["Database", "DatabaseURL"]

databasez/core/database.py

@@ -328,7 +328,7 @@ async def decr_refcount(self) -> bool:
         return False
 
     async def connect_hook(self) -> None:
-        """Refcount protected connect hook. Executed begore engine and global connection setup."""
+        """Refcount protected connect hook. Executed before engine and global connection setup."""
 
     async def connect(self) -> bool:
         """
@@ -343,9 +343,9 @@ async def connect(self) -> bool:
                 raise RuntimeError("Subdatabases and polling are disabled")
             # copy when not in map
             if loop not in self._databases_map:
-                assert (
-                    self._global_connection is not None
-                ), "global connection should have been set"
+                assert self._global_connection is not None, (
+                    "global connection should have been set"
+                )
                 # correctly initialize force_rollback with parent value
                 database = self.__class__(
                     self, force_rollback=bool(self.force_rollback), full_isolation=False
@@ -405,7 +405,9 @@ async def disconnect(
             assert not self._databases_map, "sub databases still active, force terminate them"
             for sub_database in self._databases_map.values():
                 await arun_coroutine_threadsafe(
-                    sub_database.disconnect(True), sub_database._loop, self.poll_interval
+                    sub_database.disconnect(True),
+                    sub_database._loop,
+                    self.poll_interval,
                 )
             self._databases_map = {}
         assert not self._databases_map, "sub databases still active"
@@ -556,7 +558,9 @@ async def drop_all(self, meta: MetaData, timeout: float | None = None, **kwargs:
     @multiloop_protector(False)
     def _non_global_connection(
         self,
-        timeout: float | None = None,  # stub for type checker, multiloop_protector handles timeout
+        timeout: (
+            float | None
+        ) = None,  # stub for type checker, multiloop_protector handles timeout
     ) -> Connection:
         if self._connection is None:
             _connection = self._connection = Connection(self)
@@ -624,15 +628,17 @@ def get_backends(
                 module = importlib.import_module(imp_path)
             except ImportError as exc:
                 logging.debug(
-                    f'Import of "{imp_path}" failed. This is not an error.', exc_info=exc
+                    f'Import of "{imp_path}" failed. This is not an error.',
+                    exc_info=exc,
                 )
                 if "+" in scheme:
                     imp_path = f"{overwrite_path}.{scheme.split('+', 1)[0]}"
                     try:
                         module = importlib.import_module(imp_path)
                     except ImportError as exc:
                         logging.debug(
-                            f'Import of "{imp_path}" failed. This is not an error.', exc_info=exc
+                            f'Import of "{imp_path}" failed. This is not an error.',
+                            exc_info=exc,
                         )
             if module is not None:
                 break

databasez/dialects/dbapi2.py

@@ -14,7 +14,6 @@
 from sqlalchemy.pool import AsyncAdaptedQueuePool
 from sqlalchemy.sql import text
 from sqlalchemy.util.concurrency import await_only
-from sqlalchemy_utils.functions.orm import quote
 
 from databasez.utils import AsyncWrapper
 
@@ -118,7 +117,8 @@ def has_table(
         schema: str | None = None,
         **kw: Any,
     ) -> bool:
-        stmt = text(f"select * from '{quote(connection, table_name)}' LIMIT 1")
+        quoted = self.identifier_preparer.quote(table_name)
+        stmt = text(f"select 1 from '{quoted}' LIMIT 1")
         try:
             connection.execute(stmt)
             return True

databasez/dialects/jdbc.py

@@ -27,7 +27,6 @@
 from sqlalchemy.pool import AsyncAdaptedQueuePool
 from sqlalchemy.sql import sqltypes, text
 from sqlalchemy.util.concurrency import await_only
-from sqlalchemy_utils.functions.orm import quote
 
 from databasez.utils import AsyncWrapper
 
@@ -113,7 +112,8 @@ def has_table(
         schema: str | None = None,
         **kw: Any,
     ) -> bool:
-        stmt = text(f"select * from '{quote(connection, table_name)}' LIMIT 1")
+        quoted = self.identifier_preparer.quote(table_name)
+        stmt = text(f"SELECT 1 from '{quoted}' LIMIT 1")
         try:
             connection.execute(stmt)
             return True

databasez/sqlalchemy.py

@@ -57,9 +57,9 @@ async def start(self, is_root: bool, **extra_options: Any) -> None:
             self.old_transaction_level = ""
         if extra_options:
             await connection.execution_options(**extra_options)
-        assert (
-            await connection.get_isolation_level() != "AUTOCOMMIT"
-        ), "transactions doesn't work with AUTOCOMMIT. Please specify another transaction level."
+        assert await connection.get_isolation_level() != "AUTOCOMMIT", (
+            "transactions doesn't work with AUTOCOMMIT. Please specify another transaction level."
+        )
         if in_transaction:
             self.raw_transaction = await connection.begin_nested()
         else:

databasez/testclient.py

@@ -5,19 +5,9 @@
 
 import sqlalchemy
 from sqlalchemy.exc import OperationalError, ProgrammingError
-from sqlalchemy_utils.functions.database import _sqlite_file_exists
-from sqlalchemy_utils.functions.orm import quote
 
 from databasez import Database, DatabaseURL
-from databasez.utils import DATABASEZ_POLL_INTERVAL, ThreadPassingExceptions
-
-
-async def _get_scalar_result(engine: Any, sql: Any) -> Any:
-    try:
-        async with engine.connect() as conn:
-            return await conn.scalar(sql)
-    except Exception:
-        return False
+from databasez.utils import DATABASEZ_POLL_INTERVAL, ThreadPassingExceptions, get_quoter
 
 
 class DatabaseTestClient(Database):
@@ -146,41 +136,60 @@ async def database_exists(cls, url: Union[str, "sqlalchemy.URL", DatabaseURL]) -
         database = url.database
         dialect_name = url.sqla_url.get_dialect(True).name
         if dialect_name == "postgresql":
-            text = f"SELECT 1 FROM pg_database WHERE datname='{database}'"
-            for db in (database, "postgres", "template1", "template0", None):
+            text = "SELECT 1 FROM pg_database WHERE datname=:database"
+            for db in (database, None, "postgres", "template1", "template0"):
                 url = url.replace(database=db)
-                async with Database(url, full_isolation=False, force_rollback=False) as db_client:
-                    try:
-                        return bool(
-                            await _get_scalar_result(db_client.engine, sqlalchemy.text(text))
-                        )
-                    except (ProgrammingError, OperationalError):
-                        pass
+                try:
+                    async with Database(
+                        url, full_isolation=False, force_rollback=False
+                    ) as db_client:
+                        if await db_client.fetch_val(
+                            # if we can connect to the db, it exists
+                            "SELECT 1"
+                            if db == database
+                            else sqlalchemy.text(text).bindparams(database=database)
+                        ):
+                            return True
+                except Exception:
+                    pass
             return False
 
         elif dialect_name == "mysql":
-            url = url.replace(database=None)
-            text = (
-                "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA "
-                f"WHERE SCHEMA_NAME = '{database}'"
-            )
-            async with Database(url, full_isolation=False, force_rollback=False) as db_client:
-                return bool(await _get_scalar_result(db_client.engine, sqlalchemy.text(text)))
+            for db in (database, None, "root"):
+                url = url.replace(database=db)
+                try:
+                    async with Database(
+                        url, full_isolation=False, force_rollback=False
+                    ) as db_client:
+                        if await db_client.fetch_val(
+                            (
+                                # if we can connect to the db, it exists
+                                "SELECT 1"
+                                if db == database
+                                else sqlalchemy.text(
+                                    "SELECT 1 FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = :database"
+                                ).bindparams(database=database)
+                            ),
+                        ):
+                            return True
+                except Exception:
+                    pass
+            return False
 
         elif dialect_name == "sqlite":
             if database:
-                return database == ":memory:" or _sqlite_file_exists(database)
+                return database == ":memory:" or os.path.exists(database)
             else:
                 # The default SQLAlchemy database is in memory, and :memory: is
                 # not required, thus we should support that use case.
                 return True
         else:
-            text = "SELECT 1"
-            async with Database(url, full_isolation=False, force_rollback=False) as db_client:
-                try:
-                    return bool(await _get_scalar_result(db_client.engine, sqlalchemy.text(text)))
-                except (ProgrammingError, OperationalError):
-                    return False
+            try:
+                async with Database(url, full_isolation=False, force_rollback=False) as db_client:
+                    await db_client.fetch_val("SELECT 1")
+                    return True
+            except Exception:
+                return False
 
     @classmethod
     async def create_database(
@@ -191,8 +200,9 @@ async def create_database(
     ) -> None:
         url = url if isinstance(url, DatabaseURL) else DatabaseURL(url)
         database = url.database
-        dialect_name = url.sqla_url.get_dialect(True).name
-        dialect_driver = url.sqla_url.get_dialect(True).driver
+        dialect = url.sqla_url.get_dialect(True)
+        dialect_name = dialect.name
+        dialect_driver = dialect.driver
 
         # we don't want to connect to a not existing db
         if dialect_name == "postgresql":
@@ -209,7 +219,10 @@ async def create_database(
             and dialect_driver in {"asyncpg", "pg8000", "psycopg", "psycopg2", "psycopg2cffi"}
         ):
             db_client = Database(
-                url, isolation_level="AUTOCOMMIT", force_rollback=False, full_isolation=False
+                url,
+                isolation_level="AUTOCOMMIT",
+                force_rollback=False,
+                full_isolation=False,
             )
         else:
             db_client = Database(url, force_rollback=False, full_isolation=False)
@@ -218,29 +231,34 @@ async def create_database(
                 if not template:
                     template = "template1"
 
-                async with db_client.engine.begin() as conn:  # type: ignore
+                async with db_client.connection() as conn:
+                    quote = get_quoter(conn.async_connection)
                     text = (
-                        f"CREATE DATABASE {quote(conn, database)} ENCODING "
-                        f"'{encoding}' TEMPLATE {quote(conn, template)}"
+                        f"CREATE DATABASE {quote(database)} ENCODING "
+                        f"'{encoding}' TEMPLATE {quote(template)}"
                     )
                     await conn.execute(sqlalchemy.text(text))
 
             elif dialect_name == "mysql":
-                async with db_client.engine.begin() as conn:  # type: ignore
-                    text = f"CREATE DATABASE {quote(conn, database)} CHARACTER SET = '{encoding}'"
+                async with db_client.connection() as conn:
+                    quote = get_quoter(conn.async_connection)
+                    text = f"CREATE DATABASE {quote(database)} CHARACTER SET = '{quote(encoding)}'"
                     await conn.execute(sqlalchemy.text(text))
 
             elif dialect_name == "sqlite" and database != ":memory:":
                 if database:
                     # create a sqlite file
-                    async with db_client.engine.begin() as conn:  # type: ignore
-                        await conn.execute(sqlalchemy.text("CREATE TABLE DB(id int)"))
-                        await conn.execute(sqlalchemy.text("DROP TABLE DB"))
+                    async with (
+                        db_client.connection() as conn,
+                        conn.transaction(force_rollback=False),
+                    ):
+                        await conn.execute("CREATE TABLE _dropme_DB(id int)")
+                        await conn.execute("DROP TABLE _dropme_DB")
 
             else:
-                async with db_client.engine.begin() as conn:  # type: ignore
-                    text = f"CREATE DATABASE {quote(conn, database)}"
-                    await conn.execute(sqlalchemy.text(text))
+                async with db_client.connection() as conn:
+                    quote = get_quoter(conn.async_connection)
+                    await conn.execute(sqlalchemy.text(f"CREATE DATABASE {quote(database)}"))
 
     @classmethod
     async def drop_database(cls, url: Union[str, "sqlalchemy.URL", DatabaseURL]) -> None:
@@ -264,7 +282,10 @@ async def drop_database(cls, url: Union[str, "sqlalchemy.URL", DatabaseURL]) ->
             and dialect_driver in {"asyncpg", "pg8000", "psycopg", "psycopg2", "psycopg2cffi"}
         ):
             db_client = Database(
-                url, isolation_level="AUTOCOMMIT", force_rollback=False, full_isolation=False
+                url,
+                isolation_level="AUTOCOMMIT",
+                force_rollback=False,
+                full_isolation=False,
             )
         else:
             db_client = Database(url, force_rollback=False, full_isolation=False)
@@ -274,6 +295,7 @@ async def drop_database(cls, url: Union[str, "sqlalchemy.URL", DatabaseURL]) ->
                     os.remove(database)
             elif dialect_name.startswith("postgres"):
                 async with db_client.connection() as conn:
+                    quote = get_quoter(conn.async_connection)
                     # Disconnect all users from the database we are dropping.
                     server_version_raw = (
                         await conn.fetch_val(
@@ -282,7 +304,7 @@ async def drop_database(cls, url: Union[str, "sqlalchemy.URL", DatabaseURL]) ->
                     ).split(" ")[0]
                     version = tuple(map(int, server_version_raw.split(".")))
                     pid_column = "pid" if (version >= (9, 2)) else "procpid"
-                    quoted_db = quote(conn.async_connection, database)
+                    quoted_db = quote(database)
                     text = f"""
                     SELECT pg_terminate_backend(pg_stat_activity.{pid_column})
                     FROM pg_stat_activity
@@ -297,8 +319,10 @@ async def drop_database(cls, url: Union[str, "sqlalchemy.URL", DatabaseURL]) ->
                         await conn.execute(text)
             else:
                 async with db_client.connection() as conn:
-                    text = f"DROP DATABASE {quote(conn.async_connection, database)}"
-                    await conn.execute(sqlalchemy.text(text))
+                    quote = get_quoter(conn.async_connection)
+                    text = f"DROP DATABASE {quote(database)}"
+                    with contextlib.suppress(ProgrammingError):
+                        await conn.execute(sqlalchemy.text(text))
 
     def drop_db_protected(self) -> None:
         thread = ThreadPassingExceptions(

databasez/utils.py

@@ -8,6 +8,9 @@
 from threading import Thread
 from typing import Any, TypeVar, cast
 
+from sqlalchemy.engine import Connection as SQLAConnection
+from sqlalchemy.engine import Dialect
+
 DATABASEZ_RESULT_TIMEOUT: float | None = None
 # Poll with 0.1ms, this way CPU isn't at 100%
 DATABASEZ_POLL_INTERVAL: float = 0.0001
@@ -170,9 +173,11 @@ def run(self) -> None:
             self._exc_raised = exc
 
     def join(self, timeout: float | int | None = None) -> None:
-        super().join(timeout=timeout)
-        if self._exc_raised:
-            raise self._exc_raised
+        try:
+            super().join(timeout=timeout)
+        finally:
+            if self._exc_raised is not None:
+                raise self._exc_raised
 
 
 MultiloopProtectorCallable = TypeVar("MultiloopProtectorCallable", bound=Callable)
@@ -232,3 +237,19 @@ def wrapper(
         return cast(MultiloopProtectorCallable, wrapper)
 
     return _decorator
+
+
+def get_dialect(async_conn_or_dialect: SQLAConnection | Dialect, /) -> Dialect:
+    if isinstance(async_conn_or_dialect, Dialect):
+        return async_conn_or_dialect
+    if hasattr(async_conn_or_dialect, "bind"):
+        async_conn_or_dialect = async_conn_or_dialect.bind
+    return cast("SQLAConnection", async_conn_or_dialect).dialect
+
+
+def get_quoter(async_conn_or_dialect: SQLAConnection | Dialect, /) -> Callable[[str], str]:
+    # needs underlying async connection as object or dialect
+    dialect = get_dialect(async_conn_or_dialect)
+    if hasattr(dialect, "identifier_preparer"):
+        return dialect.identifier_preparer.quote
+    return dialect.preparer(dialect).quote