fix: Refactor ADK store SQL table creation methods to be asynchronous (#144)

* fix: Refactor ADK store SQL table creation methods to be asynchronous

This is done mainly to address issues with Oracle's JSON detection and in-memory flag on creation.

- Changed `_get_create_sessions_table_sql` and `_get_create_events_table_sql` methods in various ADK store classes to be asynchronous.
- Updated `create_tables` methods to use `execute_script` for executing SQL statements asynchronously.
- Modified the Oracle ADK store to support INMEMORY PRIORITY HIGH clause for table creation.
- Added tests to verify the correct creation of tables with INMEMORY settings for both async and sync stores.
- Updated Litestar configuration to reflect changes in INMEMORY clause handling.

* fix: assertions

* refactor: switch from information_schema to pg_catalog for column and table queries in PostgreSQL adapters (#146)

Author: cofin

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.14.0"
+    rev: "v0.14.1"
     hooks:
       - id: ruff
         args: ["--fix"]

sqlspec/adapters/adbc/data_dictionary.py

@@ -300,22 +300,44 @@ def get_columns(
                 for row in result.data or []
             ]
 
+        if dialect == "postgres":
+            schema_name = schema or "public"
+            sql = """
+                SELECT
+                    a.attname::text AS column_name,
+                    pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
+                    CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
+                    pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
+                FROM pg_catalog.pg_attribute a
+                JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
+                JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
+                LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
+                WHERE c.relname = ?
+                    AND n.nspname = ?
+                    AND a.attnum > 0
+                    AND NOT a.attisdropped
+                ORDER BY a.attnum
+            """
+            result = adbc_driver.execute(sql, (table, schema_name))
+            return result.data or []
+
         if schema:
-            sql = f"""
+            sql = """
                 SELECT column_name, data_type, is_nullable, column_default
                 FROM information_schema.columns
-                WHERE table_name = '{table}' AND table_schema = '{schema}'
+                WHERE table_name = ? AND table_schema = ?
                 ORDER BY ordinal_position
             """
+            result = adbc_driver.execute(sql, (table, schema))
         else:
-            sql = f"""
+            sql = """
                 SELECT column_name, data_type, is_nullable, column_default
                 FROM information_schema.columns
-                WHERE table_name = '{table}'
+                WHERE table_name = ?
                 ORDER BY ordinal_position
             """
+            result = adbc_driver.execute(sql, (table,))
 
-        result = adbc_driver.execute(sql)
         return result.data or []
 
     def list_available_features(self) -> "list[str]":

sqlspec/adapters/aiosqlite/adk/store.py

@@ -136,7 +136,7 @@ def __init__(self, config: "AiosqliteConfig") -> None:
         """
         super().__init__(config)
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get SQLite CREATE TABLE SQL for sessions.
 
         Returns:
@@ -163,7 +163,7 @@ def _get_create_sessions_table_sql(self) -> str:
             ON {self._session_table}(update_time DESC);
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get SQLite CREATE TABLE SQL for events.
 
         Returns:
@@ -228,11 +228,10 @@ async def _enable_foreign_keys(self, connection: Any) -> None:
 
     async def create_tables(self) -> None:
         """Create both sessions and events tables if they don't exist."""
-        async with self._config.provide_connection() as conn:
-            await self._enable_foreign_keys(conn)
-            await conn.executescript(self._get_create_sessions_table_sql())
-            await conn.executescript(self._get_create_events_table_sql())
-            await conn.commit()
+        async with self._config.provide_session() as driver:
+            await self._enable_foreign_keys(driver.connection)
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_session(

sqlspec/adapters/asyncmy/adk/store.py

@@ -106,7 +106,7 @@ def _parse_owner_id_column_for_mysql(self, column_ddl: str) -> "tuple[str, str]"
 
         return (col_def, fk_constraint)
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get MySQL CREATE TABLE SQL for sessions.
 
         Returns:
@@ -145,7 +145,7 @@ def _get_create_sessions_table_sql(self) -> str:
         ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get MySQL CREATE TABLE SQL for events.
 
         Returns:
@@ -199,9 +199,9 @@ def _get_drop_tables_sql(self) -> "list[str]":
 
     async def create_tables(self) -> None:
         """Create both sessions and events tables if they don't exist."""
-        async with self._config.provide_connection() as conn, conn.cursor() as cursor:
-            await cursor.execute(self._get_create_sessions_table_sql())
-            await cursor.execute(self._get_create_events_table_sql())
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_session(

sqlspec/adapters/asyncpg/adk/store.py

@@ -84,7 +84,7 @@ def __init__(self, config: AsyncConfigT) -> None:
         """
         super().__init__(config)
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get PostgreSQL CREATE TABLE SQL for sessions.
 
         Returns:
@@ -125,7 +125,7 @@ def _get_create_sessions_table_sql(self) -> str:
             WHERE state != '{{}}'::jsonb;
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get PostgreSQL CREATE TABLE SQL for events.
 
         Returns:
@@ -181,9 +181,9 @@ def _get_drop_tables_sql(self) -> "list[str]":
 
     async def create_tables(self) -> None:
         """Create both sessions and events tables if they don't exist."""
-        async with self.config.provide_connection() as conn:
-            await conn.execute(self._get_create_sessions_table_sql())
-            await conn.execute(self._get_create_events_table_sql())
+        async with self.config.provide_session() as driver:
+            await driver.execute_script(await self._get_create_sessions_table_sql())
+            await driver.execute_script(await self._get_create_events_table_sql())
         logger.debug("Created ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_session(

sqlspec/adapters/asyncpg/data_dictionary.py

@@ -117,7 +117,7 @@ async def get_optimal_type(self, driver: AsyncDriverAdapterBase, type_category:
     async def get_columns(
         self, driver: AsyncDriverAdapterBase, table: str, schema: "str | None" = None
     ) -> "list[dict[str, Any]]":
-        """Get column information for a table using information_schema.
+        """Get column information for a table using pg_catalog.
 
         Args:
             driver: AsyncPG driver instance
@@ -130,25 +130,32 @@ async def get_columns(
                 - data_type: PostgreSQL data type
                 - is_nullable: Whether column allows NULL (YES/NO)
                 - column_default: Default value if any
+
+        Notes:
+            Uses pg_catalog instead of information_schema to avoid potential
+            issues with PostgreSQL 'name' type in some drivers.
         """
         asyncpg_driver = cast("AsyncpgDriver", driver)
 
-        if schema:
-            sql = f"""
-                SELECT column_name, data_type, is_nullable, column_default
-                FROM information_schema.columns
-                WHERE table_name = '{table}' AND table_schema = '{schema}'
-                ORDER BY ordinal_position
-            """
-        else:
-            sql = f"""
-                SELECT column_name, data_type, is_nullable, column_default
-                FROM information_schema.columns
-                WHERE table_name = '{table}' AND table_schema = 'public'
-                ORDER BY ordinal_position
-            """
-
-        result = await asyncpg_driver.execute(sql)
+        schema_name = schema or "public"
+        sql = """
+            SELECT
+                a.attname::text AS column_name,
+                pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
+                CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END AS is_nullable,
+                pg_catalog.pg_get_expr(d.adbin, d.adrelid)::text AS column_default
+            FROM pg_catalog.pg_attribute a
+            JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
+            JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid
+            LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
+            WHERE c.relname = $1
+                AND n.nspname = $2
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+            ORDER BY a.attnum
+        """
+
+        result = await asyncpg_driver.execute(sql, (table, schema_name))
         return result.data or []
 
     def list_available_features(self) -> "list[str]":

sqlspec/adapters/bigquery/adk/store.py

@@ -8,7 +8,7 @@
 from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord
 from sqlspec.utils.logging import get_logger
 from sqlspec.utils.serializers import from_json, to_json
-from sqlspec.utils.sync_tools import async_
+from sqlspec.utils.sync_tools import async_, run_
 
 if TYPE_CHECKING:
     from sqlspec.adapters.bigquery.config import BigQueryConfig
@@ -102,7 +102,7 @@ def _get_full_table_name(self, table_name: str) -> str:
             return f"`{self._dataset_id}.{table_name}`"
         return f"`{table_name}`"
 
-    def _get_create_sessions_table_sql(self) -> str:
+    async def _get_create_sessions_table_sql(self) -> str:
         """Get BigQuery CREATE TABLE SQL for sessions.
 
         Returns:
@@ -136,7 +136,7 @@ def _get_create_sessions_table_sql(self) -> str:
         CLUSTER BY app_name, user_id
         """
 
-    def _get_create_events_table_sql(self) -> str:
+    async def _get_create_events_table_sql(self) -> str:
         """Get BigQuery CREATE TABLE SQL for events.
 
         Returns:
@@ -193,9 +193,9 @@ def _get_drop_tables_sql(self) -> "list[str]":
 
     def _create_tables(self) -> None:
         """Synchronous implementation of create_tables."""
-        with self._config.provide_connection() as conn:
-            conn.query(self._get_create_sessions_table_sql()).result()
-            conn.query(self._get_create_events_table_sql()).result()
+        with self._config.provide_session() as driver:
+            driver.execute_script(run_(self._get_create_sessions_table_sql)())
+            driver.execute_script(run_(self._get_create_events_table_sql)())
         logger.debug("Created BigQuery ADK tables: %s, %s", self._session_table, self._events_table)
 
     async def create_tables(self) -> None: