llama-stack 0.3.5__py3-none-any.whl → 0.4.1__py3-none-any.whl

llama_stack/{providers/utils → core/storage}/kvstore/mongodb/mongodb.py

@@ -9,8 +9,8 @@ from datetime import datetime
 from pymongo import AsyncMongoClient
 from pymongo.asynchronous.collection import AsyncCollection
 
+from llama_stack.core.storage.kvstore import KVStore
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import KVStore
 
 from ..config import MongoDBKVStoreConfig
 
@@ -30,14 +30,13 @@ class MongoDBKVStoreImpl(KVStore):
 
     async def initialize(self) -> None:
         try:
-            conn_creds = {
-                "host": self.config.host,
-                "port": self.config.port,
-                "username": self.config.user,
-                "password": self.config.password,
-            }
-            conn_creds = {k: v for k, v in conn_creds.items() if v is not None}
-            self.conn = AsyncMongoClient(**conn_creds)
+            # Pass parameters explicitly to satisfy mypy - AsyncMongoClient doesn't accept **dict
+            self.conn = AsyncMongoClient(
+                host=self.config.host if self.config.host is not None else None,
+                port=self.config.port if self.config.port is not None else None,
+                username=self.config.user if self.config.user is not None else None,
+                password=self.config.password if self.config.password is not None else None,
+            )
         except Exception as e:
             log.exception("Could not connect to MongoDB database server")
             raise RuntimeError("Could not connect to MongoDB database server") from e
@@ -79,4 +78,8 @@ class MongoDBKVStoreImpl(KVStore):
         end_key = self._namespaced_key(end_key)
         query = {"key": {"$gte": start_key, "$lt": end_key}}
         cursor = self.collection.find(query, {"key": 1, "_id": 0}).sort("key", 1)
-        return [doc["key"] for doc in cursor]
+        # AsyncCursor requires async iteration
+        result = []
+        async for doc in cursor:
+            result.append(doc["key"])
+        return result

llama_stack/{providers/utils → core/storage}/kvstore/postgres/postgres.py

@@ -6,12 +6,13 @@
 
 from datetime import datetime
 
-import psycopg2
-from psycopg2.extras import DictCursor
+import psycopg2  # type: ignore[import-not-found]
+from psycopg2.extensions import connection as PGConnection  # type: ignore[import-not-found]
+from psycopg2.extras import DictCursor  # type: ignore[import-not-found]
 
 from llama_stack.log import get_logger
+from llama_stack_api.internal.kvstore import KVStore
 
-from ..api import KVStore
 from ..config import PostgresKVStoreConfig
 
 log = get_logger(name=__name__, category="providers::utils")
@@ -20,12 +21,12 @@ log = get_logger(name=__name__, category="providers::utils")
 class PostgresKVStoreImpl(KVStore):
     def __init__(self, config: PostgresKVStoreConfig):
         self.config = config
-        self.conn = None
-        self.cursor = None
+        self._conn: PGConnection | None = None
+        self._cursor: DictCursor | None = None
 
     async def initialize(self) -> None:
         try:
-            self.conn = psycopg2.connect(
+            self._conn = psycopg2.connect(
                 host=self.config.host,
                 port=self.config.port,
                 database=self.config.db,
@@ -34,11 +35,11 @@ class PostgresKVStoreImpl(KVStore):
                 sslmode=self.config.ssl_mode,
                 sslrootcert=self.config.ca_cert_path,
             )
-            self.conn.autocommit = True
-            self.cursor = self.conn.cursor(cursor_factory=DictCursor)
+            self._conn.autocommit = True
+            self._cursor = self._conn.cursor(cursor_factory=DictCursor)
 
             # Create table if it doesn't exist
-            self.cursor.execute(
+            self._cursor.execute(
                 f"""
                 CREATE TABLE IF NOT EXISTS {self.config.table_name} (
                     key TEXT PRIMARY KEY,
@@ -51,6 +52,11 @@ class PostgresKVStoreImpl(KVStore):
             log.exception("Could not connect to PostgreSQL database server")
             raise RuntimeError("Could not connect to PostgreSQL database server") from e
 
+    def _cursor_or_raise(self) -> DictCursor:
+        if self._cursor is None:
+            raise RuntimeError("Postgres client not initialized")
+        return self._cursor
+
     def _namespaced_key(self, key: str) -> str:
         if not self.config.namespace:
             return key
@@ -58,7 +64,8 @@ class PostgresKVStoreImpl(KVStore):
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         key = self._namespaced_key(key)
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"""
             INSERT INTO {self.config.table_name} (key, value, expiration)
             VALUES (%s, %s, %s)
@@ -70,7 +77,8 @@ class PostgresKVStoreImpl(KVStore):
 
     async def get(self, key: str) -> str | None:
         key = self._namespaced_key(key)
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"""
             SELECT value FROM {self.config.table_name}
             WHERE key = %s
@@ -78,12 +86,13 @@ class PostgresKVStoreImpl(KVStore):
             """,
             (key,),
         )
-        result = self.cursor.fetchone()
+        result = cursor.fetchone()
         return result[0] if result else None
 
     async def delete(self, key: str) -> None:
         key = self._namespaced_key(key)
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"DELETE FROM {self.config.table_name} WHERE key = %s",
             (key,),
         )
@@ -92,7 +101,8 @@ class PostgresKVStoreImpl(KVStore):
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
 
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"""
             SELECT value FROM {self.config.table_name}
             WHERE key >= %s AND key < %s
@@ -101,14 +111,15 @@ class PostgresKVStoreImpl(KVStore):
             """,
             (start_key, end_key),
         )
-        return [row[0] for row in self.cursor.fetchall()]
+        return [row[0] for row in cursor.fetchall()]
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
 
-        self.cursor.execute(
+        cursor = self._cursor_or_raise()
+        cursor.execute(
             f"SELECT key FROM {self.config.table_name} WHERE key >= %s AND key < %s",
             (start_key, end_key),
         )
-        return [row[0] for row in self.cursor.fetchall()]
+        return [row[0] for row in cursor.fetchall()]

llama_stack/{providers/utils → core/storage}/kvstore/redis/redis.py

@@ -6,18 +6,25 @@
 
 from datetime import datetime
 
-from redis.asyncio import Redis
+from redis.asyncio import Redis  # type: ignore[import-not-found]
+
+from llama_stack_api.internal.kvstore import KVStore
 
-from ..api import KVStore
 from ..config import RedisKVStoreConfig
 
 
 class RedisKVStoreImpl(KVStore):
     def __init__(self, config: RedisKVStoreConfig):
         self.config = config
+        self._redis: Redis | None = None
 
     async def initialize(self) -> None:
-        self.redis = Redis.from_url(self.config.url)
+        self._redis = Redis.from_url(self.config.url)
+
+    def _client(self) -> Redis:
+        if self._redis is None:
+            raise RuntimeError("Redis client not initialized")
+        return self._redis
 
     def _namespaced_key(self, key: str) -> str:
         if not self.config.namespace:
@@ -26,30 +33,37 @@ class RedisKVStoreImpl(KVStore):
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         key = self._namespaced_key(key)
-        await self.redis.set(key, value)
+        client = self._client()
+        await client.set(key, value)
         if expiration:
-            await self.redis.expireat(key, expiration)
+            await client.expireat(key, expiration)
 
     async def get(self, key: str) -> str | None:
         key = self._namespaced_key(key)
-        value = await self.redis.get(key)
+        client = self._client()
+        value = await client.get(key)
         if value is None:
             return None
-        await self.redis.ttl(key)
-        return value
+        await client.ttl(key)
+        if isinstance(value, bytes):
+            return value.decode("utf-8")
+        if isinstance(value, str):
+            return value
+        return str(value)
 
     async def delete(self, key: str) -> None:
         key = self._namespaced_key(key)
-        await self.redis.delete(key)
+        await self._client().delete(key)
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
+        client = self._client()
         cursor = 0
         pattern = start_key + "*"  # Match all keys starting with start_key prefix
-        matching_keys = []
+        matching_keys: list[str | bytes] = []
         while True:
-            cursor, keys = await self.redis.scan(cursor, match=pattern, count=1000)
+            cursor, keys = await client.scan(cursor, match=pattern, count=1000)
 
             for key in keys:
                 key_str = key.decode("utf-8") if isinstance(key, bytes) else key
@@ -61,7 +75,7 @@ class RedisKVStoreImpl(KVStore):
 
         # Then fetch all values in a single MGET call
         if matching_keys:
-            values = await self.redis.mget(matching_keys)
+            values = await client.mget(matching_keys)
             return [
                 value.decode("utf-8") if isinstance(value, bytes) else value for value in values if value is not None
             ]
@@ -70,7 +84,18 @@ class RedisKVStoreImpl(KVStore):
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         """Get all keys in the given range."""
-        matching_keys = await self.redis.zrangebylex(self.namespace, f"[{start_key}", f"[{end_key}")
-        if not matching_keys:
-            return []
-        return [k.decode("utf-8") for k in matching_keys]
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        client = self._client()
+        cursor = 0
+        pattern = start_key + "*"
+        result: list[str] = []
+        while True:
+            cursor, keys = await client.scan(cursor, match=pattern, count=1000)
+            for key in keys:
+                key_str = key.decode("utf-8") if isinstance(key, bytes) else str(key)
+                if start_key <= key_str <= end_key:
+                    result.append(key_str)
+            if cursor == 0:
+                break
+        return result

llama_stack/{providers/utils → core/storage}/kvstore/sqlite/sqlite.py

@@ -10,8 +10,8 @@ from datetime import datetime
 import aiosqlite
 
 from llama_stack.log import get_logger
+from llama_stack_api.internal.kvstore import KVStore
 
-from ..api import KVStore
 from ..config import SqliteKVStoreConfig
 
 logger = get_logger(name=__name__, category="providers::utils")

llama_stack/core/storage/sqlstore/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from llama_stack_api.internal.sqlstore import (
+    ColumnDefinition as ColumnDefinition,
+)
+from llama_stack_api.internal.sqlstore import (
+    ColumnType as ColumnType,
+)
+from llama_stack_api.internal.sqlstore import (
+    SqlStore as SqlStore,
+)
+
+from .sqlstore import *  # noqa: F401,F403

llama_stack/{providers/utils → core/storage}/sqlstore/authorized_sqlstore.py

@@ -14,8 +14,8 @@ from llama_stack.core.datatypes import User
 from llama_stack.core.request_headers import get_authenticated_user
 from llama_stack.core.storage.datatypes import StorageBackendType
 from llama_stack.log import get_logger
-
-from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
+from llama_stack_api import PaginatedResponse
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType, SqlStore
 
 logger = get_logger(name=__name__, category="providers::utils")
 
@@ -23,17 +23,26 @@ logger = get_logger(name=__name__, category="providers::utils")
 # WARNING: If default_policy() changes, this constant must be updated accordingly
 # or SQL filtering will fall back to conservative mode (safe but less performant)
 #
-# This policy represents: "Permit all actions when user is in owners list for ALL attribute categories"
+# This policy represents: "Permit all actions when user is in owners list for ANY attribute category"
 # The corresponding SQL logic is implemented in _build_default_policy_where_clause():
 # - Public records (no access_attributes) are always accessible
-# - Records with access_attributes require user to match ALL categories that exist in the resource
-# - Missing categories in the resource are treated as "no restriction" (allow)
+# - Records with access_attributes require user to match ANY category that exists in the resource
 # - Within each category, user needs ANY matching value (OR logic)
-# - Between categories, user needs ALL categories to match (AND logic)
+# - Between categories, user needs ANY category to match (OR logic)
 SQL_OPTIMIZED_POLICY = [
     AccessRule(
         permit=Scope(actions=list(Action)),
-        when=["user in owners roles", "user in owners teams", "user in owners projects", "user in owners namespaces"],
+        when=["user in owners " + name],
+    )
+    for name in ["roles", "teams", "projects", "namespaces"]
+] + [
+    AccessRule(
+        permit=Scope(actions=list(Action)),
+        when=["user is owner"],
+    ),
+    AccessRule(
+        permit=Scope(actions=list(Action)),
+        when=["resource is unowned"],
     ),
 ]
 
@@ -56,7 +65,7 @@ def _enhance_item_with_access_control(item: Mapping[str, Any], current_user: Use
 
 
 class SqlRecord(ProtectedResource):
-    def __init__(self, record_id: str, table_name: str, owner: User):
+    def __init__(self, record_id: str, table_name: str, owner: User | None):
         self.type = f"sql_record::{table_name}"
         self.identifier = record_id
         self.owner = owner
@@ -129,6 +138,23 @@ class AuthorizedSqlStore:
             enhanced_data = [_enhance_item_with_access_control(item, current_user) for item in data]
         await self.sql_store.insert(table, enhanced_data)
 
+    async def upsert(
+        self,
+        table: str,
+        data: Mapping[str, Any],
+        conflict_columns: list[str],
+        update_columns: list[str] | None = None,
+    ) -> None:
+        """Upsert a row with automatic access control attribute capture."""
+        current_user = get_authenticated_user()
+        enhanced_data = _enhance_item_with_access_control(data, current_user)
+        await self.sql_store.upsert(
+            table=table,
+            data=enhanced_data,
+            conflict_columns=conflict_columns,
+            update_columns=update_columns,
+        )
+
     async def fetch_all(
         self,
         table: str,
@@ -136,6 +162,7 @@ class AuthorizedSqlStore:
         limit: int | None = None,
         order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
         cursor: tuple[str, str] | None = None,
+        action: Action = Action.READ,
     ) -> PaginatedResponse:
         """Fetch all rows with automatic access control filtering."""
         access_where = self._build_access_control_where_clause(self.policy)
@@ -153,14 +180,18 @@ class AuthorizedSqlStore:
 
         for row in rows.data:
             stored_access_attrs = row.get("access_attributes")
-            stored_owner_principal = row.get("owner_principal") or ""
+            stored_owner_principal = row.get("owner_principal")
 
             record_id = row.get("id", "unknown")
-            sql_record = SqlRecord(
-                str(record_id), table, User(principal=stored_owner_principal, attributes=stored_access_attrs)
+            # Create owner as None if owner_principal is empty/missing, matching ResourceWithOwner behavior
+            owner = (
+                User(principal=stored_owner_principal, attributes=stored_access_attrs)
+                if stored_owner_principal
+                else None
             )
+            sql_record = SqlRecord(str(record_id), table, owner)
 
-            if is_action_allowed(self.policy, Action.READ, sql_record, current_user):
+            if is_action_allowed(self.policy, action, sql_record, current_user):
                 filtered_rows.append(row)
 
         return PaginatedResponse(
@@ -173,6 +204,7 @@ class AuthorizedSqlStore:
         table: str,
         where: Mapping[str, Any] | None = None,
         order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
+        action: Action = Action.READ,
     ) -> dict[str, Any] | None:
         """Fetch one row with automatic access control checking."""
         results = await self.fetch_all(
@@ -180,6 +212,7 @@ class AuthorizedSqlStore:
             where=where,
             limit=1,
             order_by=order_by,
+            action=action,
         )
 
         return results.data[0] if results.data else None
@@ -255,53 +288,40 @@ class AuthorizedSqlStore:
 
         Public records are those with:
         - owner_principal = '' (empty string)
-        - access_attributes is either SQL NULL or JSON null
 
-        Note: Different databases serialize None differently:
-        - SQLite: None → JSON null (text = 'null')
-        - Postgres: None → SQL NULL (IS NULL)
+        The policy "resource is unowned" only checks if owner_principal is empty,
+        regardless of access_attributes.
         """
-        conditions = ["owner_principal = ''"]
-        if self.database_type == StorageBackendType.SQL_POSTGRES.value:
-            # Accept both SQL NULL and JSON null for Postgres compatibility
-            # This handles both old rows (SQL NULL) and new rows (JSON null)
-            conditions.append("(access_attributes IS NULL OR access_attributes::text = 'null')")
-        elif self.database_type == StorageBackendType.SQL_SQLITE.value:
-            # SQLite serializes None as JSON null
-            conditions.append("(access_attributes IS NULL OR access_attributes = 'null')")
-        else:
-            raise ValueError(f"Unsupported database type: {self.database_type}")
-        return conditions
+        return ["owner_principal = ''"]
 
     def _build_default_policy_where_clause(self, current_user: User | None) -> str:
         """Build SQL WHERE clause for the default policy.
 
         Default policy: permit all actions when user in owners [roles, teams, projects, namespaces]
-        This means user must match ALL attribute categories that exist in the resource.
+        This means user must match ANY attribute category that exists in the resource (OR logic).
         """
         base_conditions = self._get_public_access_conditions()
-        user_attr_conditions = []
-
-        if current_user and current_user.attributes:
-            for attr_key, user_values in current_user.attributes.items():
-                if user_values:
-                    value_conditions = []
-                    for value in user_values:
-                        # Check if JSON array contains the value
-                        escaped_value = value.replace("'", "''")
-                        json_text = self._json_extract_text("access_attributes", attr_key)
-                        value_conditions.append(f"({json_text} LIKE '%\"{escaped_value}\"%')")
-
-                    if value_conditions:
-                        # Check if the category is missing (NULL)
-                        category_missing = f"{self._json_extract('access_attributes', attr_key)} IS NULL"
-                        user_matches_category = f"({' OR '.join(value_conditions)})"
-                        user_attr_conditions.append(f"({category_missing} OR {user_matches_category})")
-
-            if user_attr_conditions:
-                all_requirements_met = f"({' AND '.join(user_attr_conditions)})"
-                base_conditions.append(all_requirements_met)
 
+        if current_user:
+            # Add "user is owner" condition - user's principal matches owner_principal
+            escaped_principal = current_user.principal.replace("'", "''")
+            base_conditions.append(f"owner_principal = '{escaped_principal}'")
+
+            # Add "user in owners" conditions for attribute matching
+            if current_user.attributes:
+                for attr_key, user_values in current_user.attributes.items():
+                    if user_values:
+                        value_conditions = []
+                        for value in user_values:
+                            # Check if JSON array contains the value
+                            escaped_value = value.replace("'", "''")
+                            json_text = self._json_extract_text("access_attributes", attr_key)
+                            value_conditions.append(f"({json_text} LIKE '%\"{escaped_value}\"%')")
+
+                        if value_conditions:
+                            # User matches this category if any of their values match
+                            user_matches_category = f"({' OR '.join(value_conditions)})"
+                            base_conditions.append(user_matches_category)
         return f"({' OR '.join(base_conditions)})"
 
     def _build_conservative_where_clause(self) -> str:

llama_stack/{providers/utils → core/storage}/sqlstore/sqlalchemy_sqlstore.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 from collections.abc import Mapping, Sequence
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from sqlalchemy import (
     JSON,
@@ -26,11 +26,10 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.ext.asyncio.engine import AsyncEngine
 from sqlalchemy.sql.elements import ColumnElement
 
-from llama_stack.apis.common.responses import PaginatedResponse
 from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
 from llama_stack.log import get_logger
-
-from .api import ColumnDefinition, ColumnType, SqlStore
+from llama_stack_api import PaginatedResponse
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType, SqlStore
 
 logger = get_logger(name=__name__, category="providers::utils")
 
@@ -56,29 +55,30 @@ def _build_where_expr(column: ColumnElement, value: Any) -> ColumnElement:
             raise ValueError(f"Operator mapping must have a single operator, got: {value}")
         op, operand = next(iter(value.items()))
         if op == "==" or op == "=":
-            return column == operand
+            return cast(ColumnElement[Any], column == operand)
         if op == ">":
-            return column > operand
+            return cast(ColumnElement[Any], column > operand)
         if op == "<":
-            return column < operand
+            return cast(ColumnElement[Any], column < operand)
         if op == ">=":
-            return column >= operand
+            return cast(ColumnElement[Any], column >= operand)
         if op == "<=":
-            return column <= operand
+            return cast(ColumnElement[Any], column <= operand)
         raise ValueError(f"Unsupported operator '{op}' in where mapping")
-    return column == value
+    return cast(ColumnElement[Any], column == value)
 
 
 class SqlAlchemySqlStoreImpl(SqlStore):
     def __init__(self, config: SqlAlchemySqlStoreConfig):
         self.config = config
+        self._is_sqlite_backend = "sqlite" in self.config.engine_str
         self.async_session = async_sessionmaker(self.create_engine())
         self.metadata = MetaData()
 
     def create_engine(self) -> AsyncEngine:
         # Configure connection args for better concurrency support
         connect_args = {}
-        if "sqlite" in self.config.engine_str:
+        if self._is_sqlite_backend:
             # SQLite-specific optimizations for concurrent access
             # With WAL mode, most locks resolve in milliseconds, but allow up to 5s for edge cases
             connect_args["timeout"] = 5.0
@@ -91,7 +91,7 @@ class SqlAlchemySqlStoreImpl(SqlStore):
         )
 
         # Enable WAL mode for SQLite to support concurrent readers and writers
-        if "sqlite" in self.config.engine_str:
+        if self._is_sqlite_backend:
 
             @event.listens_for(engine.sync_engine, "connect")
             def set_sqlite_pragma(dbapi_conn, connection_record):
@@ -151,6 +151,29 @@ class SqlAlchemySqlStoreImpl(SqlStore):
             await session.execute(self.metadata.tables[table].insert(), data)
             await session.commit()
 
+    async def upsert(
+        self,
+        table: str,
+        data: Mapping[str, Any],
+        conflict_columns: list[str],
+        update_columns: list[str] | None = None,
+    ) -> None:
+        table_obj = self.metadata.tables[table]
+        dialect_insert = self._get_dialect_insert(table_obj)
+        insert_stmt = dialect_insert.values(**data)
+
+        if update_columns is None:
+            update_columns = [col for col in data.keys() if col not in conflict_columns]
+
+        update_mapping = {col: getattr(insert_stmt.excluded, col) for col in update_columns}
+        conflict_cols = [table_obj.c[col] for col in conflict_columns]
+
+        stmt = insert_stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_mapping)
+
+        async with self.async_session() as session:
+            await session.execute(stmt)
+            await session.commit()
+
     async def fetch_all(
         self,
         table: str,
@@ -240,10 +263,8 @@ class SqlAlchemySqlStoreImpl(SqlStore):
                 query = query.limit(fetch_limit)
 
             result = await session.execute(query)
-            if result.rowcount == 0:
-                rows = []
-            else:
-                rows = [dict(row._mapping) for row in result]
+            # Iterate directly - if no rows, list comprehension yields empty list
+            rows = [dict(row._mapping) for row in result]
 
             # Always return pagination result
             has_more = False
@@ -335,9 +356,18 @@ class SqlAlchemySqlStoreImpl(SqlStore):
                 add_column_sql = text(f"ALTER TABLE {table} ADD COLUMN {column_name} {compiled_type}{nullable_clause}")
 
                 await conn.execute(add_column_sql)
-
         except Exception as e:
             # If any error occurs during migration, log it but don't fail
             # The table creation will handle adding the column
             logger.error(f"Error adding column {column_name} to table {table}: {e}")
             pass
+
+    def _get_dialect_insert(self, table: Table):
+        if self._is_sqlite_backend:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            return sqlite_insert(table)
+        else:
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            return pg_insert(table)