sqlspec 0.62.0__tar.gz → 0.62.1__tar.gz

{sqlspec-0.62.0 → sqlspec-0.62.1}/.pre-commit-config.yaml

@@ -25,7 +25,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.16.2"
+    rev: "v0.16.3"
     hooks:
       - id: ruff-check
         args: ["--fix"]

{sqlspec-0.62.0 → sqlspec-0.62.1}/Makefile

@@ -113,7 +113,7 @@ release:                                           ## Bump version and create re
 	@make docs
 	@make clean
 	@make build
-	@uv run bump-my-version bump $(bump)
+	@uv run bump-my-version bump --current-version "$$(uv version --short)" $(bump)
 	@uv lock --upgrade-package sqlspec >/dev/null 2>&1
 	@echo "${OK} Release complete 🎉"
 
@@ -133,7 +133,7 @@ pre-release:                                       ## Start a pre-release: make
 	@echo "${INFO} Preparing pre-release $(version)... 🧪"
 	@make clean
 	@make build
-	@uv run bump-my-version bump --new-version $(version) pre
+	@uv run bump-my-version bump --current-version "$$(uv version --short)" --new-version $(version) pre
 	@uv lock --upgrade-package sqlspec >/dev/null 2>&1
 	@echo "${OK} Pre-release $(version) complete 🧪"
 	@echo ""

{sqlspec-0.62.0 → sqlspec-0.62.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: sqlspec
-Version: 0.62.0
+Version: 0.62.1
 Summary: SQL Experiments in Python
 Project-URL: Discord, https://discord.gg/litestar
 Project-URL: Issue, https://github.com/litestar-org/sqlspec/issues/

{sqlspec-0.62.0 → sqlspec-0.62.1}/pyproject.toml

@@ -24,7 +24,7 @@ maintainers = [{ name = "Litestar Developers", email = "hello@litestar.dev" }]
 name = "sqlspec"
 readme = "README.md"
 requires-python = ">=3.10, <4.0"
-version = "0.62.0"
+version = "0.62.1"
 
 [project.urls]
 Discord = "https://discord.gg/litestar"
@@ -331,7 +331,7 @@ opt_level = "3"   # Maximum optimization (0-3)
 allow_dirty = true
 commit = false
 commit_args = "--no-verify"
-current_version = "0.61.1"
+current_version = "0.62.1"
 ignore_missing_files = false
 ignore_missing_version = false
 message = "chore(release): bump to v{new_version}"

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/asyncmy/adk/store.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast
 import asyncmy
 from typing_extensions import NotRequired
 
+from sqlspec.adapters.asyncmy.core import resolve_rowcount
 from sqlspec.config import ADKConfig
 from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options
 from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore
@@ -309,7 +310,7 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
                 await cursor.execute(sql, tuple(params))
                 await conn.commit()
-                return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+                return resolve_rowcount(cursor)
         except asyncmy.errors.ProgrammingError as exc:  # pyright: ignore[reportAttributeAccessIssue]
             if _is_mysql_table_missing(exc):
                 return 0
@@ -327,7 +328,7 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
                 await cursor.execute(sql, tuple(params))
                 await conn.commit()
-                return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+                return resolve_rowcount(cursor)
         except asyncmy.errors.ProgrammingError as exc:  # pyright: ignore[reportAttributeAccessIssue]
             if _is_mysql_table_missing(exc):
                 return 0
@@ -345,7 +346,7 @@ class AsyncmyADKStore(BaseAsyncADKStore["AsyncmyConfig"]):
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
                 await cursor.execute(sql, tuple(params))
                 await conn.commit()
-                return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+                return resolve_rowcount(cursor)
         except asyncmy.errors.ProgrammingError as exc:  # pyright: ignore[reportAttributeAccessIssue]
             if _is_mysql_table_missing(exc):
                 return 0
@@ -545,7 +546,7 @@ class AsyncmyADKMemoryStore(BaseAsyncADKMemoryStore["AsyncmyConfig"]):
                             entry["inserted_at"],
                         )
                     await cursor.execute(sql, params)
-                    inserted_count += cursor.rowcount
+                    inserted_count += resolve_rowcount(cursor)
             await conn.commit()
         return inserted_count
 
@@ -608,7 +609,7 @@ class AsyncmyADKMemoryStore(BaseAsyncADKMemoryStore["AsyncmyConfig"]):
         async with self._config.provide_connection() as conn, conn.cursor() as cursor:
             await cursor.execute(sql, (session_id,))
             await conn.commit()
-            return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+            return resolve_rowcount(cursor)
 
     async def delete_entries_older_than(
         self, days: int, app_name: "str | None" = None, scope: "str | None" = None
@@ -632,7 +633,7 @@ class AsyncmyADKMemoryStore(BaseAsyncADKMemoryStore["AsyncmyConfig"]):
         async with self._config.provide_connection() as conn, conn.cursor() as cursor:
             await cursor.execute(sql, tuple(params))
             await conn.commit()
-            return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+            return resolve_rowcount(cursor)
 
     async def _memory_table_ddl(self) -> str:
         """Get MySQL CREATE TABLE SQL for memory entries."""

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/asyncmy/config.py

@@ -325,7 +325,7 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm
 
         Future driver_features can be added here if needed.
         """
-        return cast("AsyncmyPool", await asyncmy.create_pool(**_pool_config(self.connection_config)))
+        return await asyncmy.create_pool(**_pool_config(self.connection_config))
 
     async def _ensure_connection(self, connection: "AsyncmyConnection") -> None:
         """Ensure connection callback has been called exactly once for this connection.

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/asyncmy/litestar/store.py

@@ -3,6 +3,7 @@
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Final, cast
 
+from sqlspec.adapters.asyncmy.core import resolve_rowcount
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.extensions.litestar.store import BaseSQLSpecStore
 from sqlspec.utils.logging import get_logger
@@ -227,7 +228,7 @@ class AsyncmyStore(BaseSQLSpecStore["AsyncmyConfig"]):
         async with self._config.provide_connection() as conn, conn.cursor() as cursor:
             await cursor.execute(sql)
             await conn.commit()
-            count: int = cursor.rowcount
+            count = resolve_rowcount(cursor)
             if count > 0:
                 self._log_delete_expired(count)
             return count

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/asyncpg/adk/store.py

@@ -599,6 +599,8 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
             return
 
         async with self._config.provide_session() as driver:
+            if self._enable_bm25:
+                self._config._ensure_pg_textsearch_available()
             await driver.execute_script(await self._memory_table_ddl())
 
     async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int:
@@ -616,9 +618,9 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
                     sql = f"""
                     INSERT INTO {self._memory_table}
                     (id, session_id, app_name, user_id, scope, event_id, author,
-                     {self._owner_id_column_name}, timestamp, content_json,
+                     {self._owner_id_column_name}, timestamp, embedding, content_json,
                      content_text, metadata_json, inserted_at)
-                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::float8[]::vector, $11, $12, $13, $14)
                     ON CONFLICT (event_id) DO NOTHING
                     """
                     result = await conn.execute(
@@ -632,6 +634,7 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
                         entry["author"],
                         owner_id,
                         entry["timestamp"],
+                        entry.get("embedding"),
                         entry["content_json"],
                         entry["content_text"],
                         entry["metadata_json"],
@@ -641,8 +644,8 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
                     sql = f"""
                     INSERT INTO {self._memory_table}
                     (id, session_id, app_name, user_id, scope, event_id, author,
-                     timestamp, content_json, content_text, metadata_json, inserted_at)
-                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
+                     timestamp, embedding, content_json, content_text, metadata_json, inserted_at)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::float8[]::vector, $10, $11, $12, $13)
                     ON CONFLICT (event_id) DO NOTHING
                     """
                     result = await conn.execute(
@@ -655,13 +658,14 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
                         entry["event_id"],
                         entry["author"],
                         entry["timestamp"],
+                        entry.get("embedding"),
                         entry["content_json"],
                         entry["content_text"],
                         entry["metadata_json"],
                         entry["inserted_at"],
                     )
                 try:
-                    inserted_count += int(result.split(" ")[1])
+                    inserted_count += int(result.rsplit(" ", 1)[-1])
                 except (IndexError, ValueError):
                     continue
 
@@ -706,7 +710,7 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
             candidate_limit = max(limit_value * 2, 50)
             sql = f"""
             WITH vector_matches AS (
-                SELECT id, RANK() OVER (ORDER BY embedding <=> {p_vec}) AS rank_vec
+                SELECT id, RANK() OVER (ORDER BY embedding <=> {p_vec}::float8[]::vector) AS rank_vec
                 FROM {self._memory_table}
                 WHERE {where_scope} AND embedding IS NOT NULL
                 LIMIT {p_cand}
@@ -732,7 +736,7 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
             sql = f"""
             SELECT * FROM {self._memory_table}
             WHERE {where_scope} AND embedding IS NOT NULL
-            ORDER BY embedding <=> {p_vec} ASC, timestamp DESC
+            ORDER BY embedding <=> {p_vec}::float8[]::vector ASC, timestamp DESC
             LIMIT {p_lim}
             """
             params = (*scope_params, list(embedding), limit_value)
@@ -759,6 +763,8 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
             params = (*scope_params, f"%{query}%", limit_value)
 
         async with self._config.provide_connection() as conn:
+            if embedding is not None and self._enable_bm25 and query:
+                self._config._ensure_pg_textsearch_available()
             rows = await conn.fetch(sql, *params)
         return [cast("StoredMemory", dict(row)) for row in rows]
 
@@ -821,7 +827,7 @@ class AsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["AsyncpgConfig"]):
 
         if self._enable_bm25:
             indexes.append(
-                f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_bm25 ON {self._memory_table} USING bm25 (content_text);"
+                f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_bm25 ON {self._memory_table} USING bm25 (content_text) WITH (text_config='english');"
             )
 
         if self._vector_index_type == "scann":

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/asyncpg/config.py

@@ -326,6 +326,8 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         self._alloydb_connector: Any | None = None
         self._pgvector_available: bool | None = None
         self._paradedb_available: bool | None = None
+        self._pg_textsearch_available: bool | None = None
+        self._pg_textsearch_probe_error: Exception | None = None
 
         self._validate_connector_config()
 
@@ -452,17 +454,28 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         if self._pgvector_available is None:
             detected_extensions: set[str] = set()
             extensions = build_postgres_extension_probe_names(self.driver_features)
+            adk_config = self.extension_config.get("adk", {})
+            bm25_enabled = bool(
+                isinstance(adk_config, dict)
+                and adk_config.get("enable_memory", True)
+                and adk_config.get("enable_bm25", False)
+            )
+            if bm25_enabled:
+                extensions.append("pg_textsearch")
             if extensions:
                 try:
                     results = await connection.fetch(
                         "SELECT extname FROM pg_extension WHERE extname = ANY($1::text[])", extensions
                     )
                     detected_extensions = {r["extname"] for r in results}
-                except Exception:
+                except Exception as exc:
                     detected_extensions = set()
+                    if bm25_enabled:
+                        self._pg_textsearch_probe_error = exc
             self.statement_config, self._pgvector_available, self._paradedb_available = (
                 resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions)
             )
+            self._pg_textsearch_available = "pg_textsearch" in detected_extensions if bm25_enabled else False
 
         if self._pgvector_available:
             await register_pgvector_support(connection)
@@ -471,6 +484,12 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         if self._user_connection_hook is not None:
             await self._user_connection_hook(connection)
 
+    def _ensure_pg_textsearch_available(self) -> None:
+        if self._pg_textsearch_available:
+            return
+        msg = "ADK memory enable_bm25 requires the pg_textsearch PostgreSQL extension"
+        raise ImproperConfigurationError(msg) from self._pg_textsearch_probe_error
+
     async def _close_pool(self) -> None:
         """Close the actual async connection pool and cleanup connectors."""
         if self.connection_instance:

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/psycopg/adk/store.py

@@ -1240,6 +1240,8 @@ class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]):
             return
 
         async with self._config.provide_session() as driver:
+            if self._enable_bm25:
+                self._config._ensure_pg_textsearch_available()
             await driver.execute_script(await self._memory_table_ddl())
 
     async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int:
@@ -1256,10 +1258,10 @@ class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]):
             query = pg_sql.SQL("""
             INSERT INTO {table} (
                 id, session_id, app_name, user_id, scope, event_id, author,
-                {owner_id_col}, timestamp, content_json, content_text,
+                {owner_id_col}, timestamp, embedding, content_json, content_text,
                 metadata_json, inserted_at
             ) VALUES (
-                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::float8[]::vector, %s, %s, %s, %s
             )
             ON CONFLICT (event_id) DO NOTHING
             """).format(
@@ -1269,9 +1271,9 @@ class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]):
             query = pg_sql.SQL("""
             INSERT INTO {table} (
                 id, session_id, app_name, user_id, scope, event_id, author,
-                timestamp, content_json, content_text, metadata_json, inserted_at
+                timestamp, embedding, content_json, content_text, metadata_json, inserted_at
             ) VALUES (
-                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+                %s, %s, %s, %s, %s, %s, %s, %s, %s::float8[]::vector, %s, %s, %s, %s
             )
             ON CONFLICT (event_id) DO NOTHING
             """).format(table=pg_sql.Identifier(self._memory_table))
@@ -1395,7 +1397,7 @@ class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]):
         sql = pg_sql.SQL(
             """
             WITH vector_matches AS (
-                SELECT id, RANK() OVER (ORDER BY embedding <=> %s) AS rank_vec
+                SELECT id, RANK() OVER (ORDER BY embedding <=> %s::float8[]::vector) AS rank_vec
                 FROM {table}
                 WHERE {where_scope} AND embedding IS NOT NULL
                 LIMIT %s
@@ -1417,6 +1419,7 @@ class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]):
         ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope)
         params = (list(embedding), *scope_params, candidate_limit, query, *scope_params, candidate_limit, limit)
         async with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur:
+            self._config._ensure_pg_textsearch_available()
             await cur.execute(sql, params)
             rows = await cur.fetchall()
         return _rows_to_records(rows)
@@ -1434,7 +1437,7 @@ class PsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["PsycopgAsyncConfig"]):
             """
             SELECT * FROM {table}
             WHERE {where_scope} AND embedding IS NOT NULL
-            ORDER BY embedding <=> %s ASC, timestamp DESC
+            ORDER BY embedding <=> %s::float8[]::vector ASC, timestamp DESC
             LIMIT %s
             """
         ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope)
@@ -1504,6 +1507,8 @@ class PsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["PsycopgSyncConfig"]):
             return
 
         with self._config.provide_session() as driver:
+            if self._enable_bm25:
+                self._config._ensure_pg_textsearch_available()
             driver.execute_script(self._memory_table_ddl())
 
     def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int:
@@ -1520,10 +1525,10 @@ class PsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["PsycopgSyncConfig"]):
             query = pg_sql.SQL("""
             INSERT INTO {table} (
                 id, session_id, app_name, user_id, scope, event_id, author,
-                {owner_id_col}, timestamp, content_json, content_text,
+                {owner_id_col}, timestamp, embedding, content_json, content_text,
                 metadata_json, inserted_at
             ) VALUES (
-                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::float8[]::vector, %s, %s, %s, %s
             )
             ON CONFLICT (event_id) DO NOTHING
             """).format(
@@ -1533,9 +1538,9 @@ class PsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["PsycopgSyncConfig"]):
             query = pg_sql.SQL("""
             INSERT INTO {table} (
                 id, session_id, app_name, user_id, scope, event_id, author,
-                timestamp, content_json, content_text, metadata_json, inserted_at
+                timestamp, embedding, content_json, content_text, metadata_json, inserted_at
             ) VALUES (
-                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+                %s, %s, %s, %s, %s, %s, %s, %s, %s::float8[]::vector, %s, %s, %s, %s
             )
             ON CONFLICT (event_id) DO NOTHING
             """).format(table=pg_sql.Identifier(self._memory_table))
@@ -1656,7 +1661,7 @@ class PsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["PsycopgSyncConfig"]):
         sql = pg_sql.SQL(
             """
             WITH vector_matches AS (
-                SELECT id, RANK() OVER (ORDER BY embedding <=> %s) AS rank_vec
+                SELECT id, RANK() OVER (ORDER BY embedding <=> %s::float8[]::vector) AS rank_vec
                 FROM {table}
                 WHERE {where_scope} AND embedding IS NOT NULL
                 LIMIT %s
@@ -1678,6 +1683,7 @@ class PsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["PsycopgSyncConfig"]):
         ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope)
         params = (list(embedding), *scope_params, candidate_limit, query, *scope_params, candidate_limit, limit)
         with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur:
+            self._config._ensure_pg_textsearch_available()
             cur.execute(sql, params)
             rows = cur.fetchall()
         return _rows_to_records(rows)
@@ -1695,7 +1701,7 @@ class PsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["PsycopgSyncConfig"]):
             """
             SELECT * FROM {table}
             WHERE {where_scope} AND embedding IS NOT NULL
-            ORDER BY embedding <=> %s ASC, timestamp DESC
+            ORDER BY embedding <=> %s::float8[]::vector ASC, timestamp DESC
             LIMIT %s
             """
         ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope)
@@ -1760,6 +1766,7 @@ def _build_insert_params(entry: "StoredMemory") -> "tuple[object, ...]":
         entry["event_id"],
         entry["author"],
         entry["timestamp"],
+        entry.get("embedding"),
         Jsonb(entry["content_json"]),
         entry["content_text"],
         Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None,
@@ -1778,6 +1785,7 @@ def _build_insert_params_with_owner(entry: "StoredMemory", owner_id: "object | N
         entry["author"],
         owner_id,
         entry["timestamp"],
+        entry.get("embedding"),
         Jsonb(entry["content_json"]),
         entry["content_text"],
         Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None,
@@ -1868,7 +1876,7 @@ def _postgres_memory_indexes(
 
     if enable_bm25:
         indexes.append(
-            f"CREATE INDEX IF NOT EXISTS idx_{memory_table}_bm25 ON {memory_table} USING bm25 (content_text);"
+            f"CREATE INDEX IF NOT EXISTS idx_{memory_table}_bm25 ON {memory_table} USING bm25 (content_text) WITH (text_config='english');"
         )
 
     if vector_index_type == "scann":

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/adapters/psycopg/config.py

@@ -312,6 +312,8 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         features_dict.setdefault("alloydb_ip_type", "PRIVATE")
         self._pgvector_available: bool | None = None
         self._paradedb_available: bool | None = None
+        self._pg_textsearch_available: bool | None = None
+        self._pg_textsearch_probe_error: Exception | None = None
         self._alloydb_connector: Any | None = None
 
         super().__init__(
@@ -428,6 +430,14 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         if self._pgvector_available is None:
             detected_extensions: set[str] = set()
             extensions = build_postgres_extension_probe_names(self.driver_features)
+            adk_config = self.extension_config.get("adk", {})
+            bm25_enabled = bool(
+                isinstance(adk_config, dict)
+                and adk_config.get("enable_memory", True)
+                and adk_config.get("enable_bm25", False)
+            )
+            if bm25_enabled:
+                extensions.append("pg_textsearch")
             if extensions:
                 try:
                     cursor = conn.execute(
@@ -435,11 +445,14 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
                     )
                     results = cursor.fetchall()
                     detected_extensions = {r[0] for r in results}  # type: ignore[index]
-                except Exception:
+                except Exception as exc:
                     detected_extensions = set()
+                    if bm25_enabled:
+                        self._pg_textsearch_probe_error = exc
             self.statement_config, self._pgvector_available, self._paradedb_available = (
                 resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions)
             )
+            self._pg_textsearch_available = "pg_textsearch" in detected_extensions if bm25_enabled else False
 
         if self._pgvector_available:
             register_pgvector_sync(conn)
@@ -452,6 +465,12 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         if self._user_connection_hook is not None:
             self._user_connection_hook(conn)
 
+    def _ensure_pg_textsearch_available(self) -> None:
+        if self._pg_textsearch_available:
+            return
+        msg = "ADK memory enable_bm25 requires the pg_textsearch PostgreSQL extension"
+        raise ImproperConfigurationError(msg) from self._pg_textsearch_probe_error
+
     def _close_pool(self) -> None:
         """Close the actual connection pool and cleanup connectors."""
         try:
@@ -638,6 +657,8 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         )
         self._pgvector_available: bool | None = None
         self._paradedb_available: bool | None = None
+        self._pg_textsearch_available: bool | None = None
+        self._pg_textsearch_probe_error: Exception | None = None
 
         super().__init__(
             connection_config=connection_config,
@@ -700,6 +721,14 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         if self._pgvector_available is None:
             detected_extensions: set[str] = set()
             extensions = build_postgres_extension_probe_names(self.driver_features)
+            adk_config = self.extension_config.get("adk", {})
+            bm25_enabled = bool(
+                isinstance(adk_config, dict)
+                and adk_config.get("enable_memory", True)
+                and adk_config.get("enable_bm25", False)
+            )
+            if bm25_enabled:
+                extensions.append("pg_textsearch")
             if extensions:
                 try:
                     cursor = await conn.execute(
@@ -707,11 +736,14 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
                     )
                     results = await cursor.fetchall()
                     detected_extensions = {r[0] for r in results}  # type: ignore[index]
-                except Exception:
+                except Exception as exc:
                     detected_extensions = set()
+                    if bm25_enabled:
+                        self._pg_textsearch_probe_error = exc
             self.statement_config, self._pgvector_available, self._paradedb_available = (
                 resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions)
             )
+            self._pg_textsearch_available = "pg_textsearch" in detected_extensions if bm25_enabled else False
 
         if self._pgvector_available:
             await register_pgvector_async(conn)
@@ -724,6 +756,12 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         if self._user_connection_hook is not None:
             await self._user_connection_hook(conn)
 
+    def _ensure_pg_textsearch_available(self) -> None:
+        if self._pg_textsearch_available:
+            return
+        msg = "ADK memory enable_bm25 requires the pg_textsearch PostgreSQL extension"
+        raise ImproperConfigurationError(msg) from self._pg_textsearch_probe_error
+
     async def _close_pool(self) -> None:
         """Close the actual async connection pool."""
         if not self.connection_instance:

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/config.py

@@ -739,7 +739,7 @@ class ADKConfig(TypedDict):
     """Dimensionality of embedding vectors (e.g. 768 for gemini-embedding-001 with MRL). Default: 768."""
 
     enable_bm25: NotRequired[bool]
-    """Enable native pg_textsearch BM25 full-text indexing on AlloyDB / PostgreSQL 17+. Default: False."""
+    """Enable native BM25 full-text indexing. Requires the pg_textsearch extension. Default: False."""
 
     scann_num_leaves: NotRequired[int]
     """Number of partition leaves (clusters) for ScaNN tree quantization. Default: 100."""
@@ -1524,8 +1524,8 @@ class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
         self.extension_config = extension_config or {}
         self.migration_config = migration_config or {}
         self._init_observability(observability_config)
-        self._initialize_migration_components()
         self.statement_config = statement_config or build_default_statement_config(default_dialect)
+        self._initialize_migration_components()
         self._storage_capabilities = None
         self.driver_features = seed_runtime_driver_features(driver_features, self.storage_capabilities())
         self._attach_lifecycle_hooks()

{sqlspec-0.62.0 → sqlspec-0.62.1}/sqlspec/extensions/adk/migrations/0001_create_adk_tables.py

@@ -8,8 +8,8 @@ Creates the canonical ADK tables:
 - adk_internal_metadata
 - adk_memory (when memory is enabled)
 
-On PostgreSQL the migration installs the ``vector`` extension before the first
-memory statement that declares a ``VECTOR`` column.
+On PostgreSQL the migration installs required extensions before the first
+memory statement that uses them.
 """
 
 import inspect
@@ -38,9 +38,11 @@ __all__ = ("down", "up")
 logger = get_logger("sqlspec.migrations.adk.create")
 
 CREATE_VECTOR_EXTENSION = "CREATE EXTENSION IF NOT EXISTS vector"
+CREATE_PG_TEXTSEARCH_EXTENSION = "CREATE EXTENSION IF NOT EXISTS pg_textsearch"
 
 _POSTGRES_DIALECTS = frozenset({"postgres", "postgresql", "pgvector", "paradedb"})
 _VECTOR_COLUMN_PATTERN = re.compile(r"\bVECTOR\s*\(", re.IGNORECASE)
+_BM25_INDEX_PATTERN = re.compile(r"\bUSING\s+bm25\s*\(", re.IGNORECASE)
 
 
 async def up(context: "MigrationContext | None" = None) -> "list[str]":
@@ -70,7 +72,7 @@ async def up(context: "MigrationContext | None" = None) -> "list[str]":
             if inspect.isawaitable(memory_sql):
                 memory_sql = await memory_sql
             memory_statements = list(memory_sql) if isinstance(memory_sql, list) else [memory_sql]
-            statements.extend(_with_vector_extension(memory_statements, context))
+            statements.extend(_with_postgres_extensions(memory_statements, context))
             log_with_context(
                 logger, logging.DEBUG, "adk.migration.create.memory.create", table_name=memory_store.memory_table
             )
@@ -98,15 +100,24 @@ async def down(context: "MigrationContext | None" = None) -> "list[str]":
     return statements
 
 
-def _with_vector_extension(statements: "list[str]", context: "MigrationContext") -> "list[str]":
-    """Prepend the PostgreSQL vector extension before the first vector statement."""
+def _with_postgres_extensions(statements: "list[str]", context: "MigrationContext") -> "list[str]":
+    """Prepend required PostgreSQL extensions before their first use."""
     dialect = (context.dialect or "").lower()
     if dialect not in _POSTGRES_DIALECTS:
         return statements
-    for index, statement in enumerate(statements):
-        if _VECTOR_COLUMN_PATTERN.search(statement):
-            return [*statements[:index], CREATE_VECTOR_EXTENSION, *statements[index:]]
-    return statements
+
+    result: list[str] = []
+    vector_enabled = False
+    pg_textsearch_enabled = False
+    for statement in statements:
+        if not vector_enabled and _VECTOR_COLUMN_PATTERN.search(statement):
+            result.append(CREATE_VECTOR_EXTENSION)
+            vector_enabled = True
+        if not pg_textsearch_enabled and _BM25_INDEX_PATTERN.search(statement):
+            result.append(CREATE_PG_TEXTSEARCH_EXTENSION)
+            pg_textsearch_enabled = True
+        result.append(statement)
+    return result
 
 
 def _raise_missing_config() -> NoReturn:

{sqlspec-0.62.0 → sqlspec-0.62.1}/tests/integration/adapters/_shared/behaviors.py

@@ -858,7 +858,7 @@ class _CursorCloseRecordingConnection:
     def cursor(self, *args: Any, **kwargs: Any) -> Any:
         closes = object.__getattribute__(self, "_closes")
         result = object.__getattribute__(self, "_inner").cursor(*args, **kwargs)
-        if inspect.isawaitable(result):
+        if inspect.isawaitable(result) and not hasattr(result, "execute"):
             return _await_recording_cursor(result, closes)
         return _CursorCloseRecordingCursor(result, closes)