agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/db/migrations/manager.py

@@ -0,0 +1,199 @@
+import importlib
+from typing import Optional, Union
+
+from packaging import version as packaging_version
+from packaging.version import Version
+
+from agno.db.base import AsyncBaseDb, BaseDb
+from agno.utils.log import log_error, log_info, log_warning
+
+
+class MigrationManager:
+    """Manager class to handle database migrations"""
+
+    available_versions: list[tuple[str, Version]] = [
+        ("v2_0_0", packaging_version.parse("2.0.0")),
+        ("v2_3_0", packaging_version.parse("2.3.0")),
+    ]
+
+    def __init__(self, db: Union[AsyncBaseDb, BaseDb]):
+        self.db = db
+
+    @property
+    def latest_schema_version(self) -> Version:
+        return self.available_versions[-1][1]
+
+    async def up(self, target_version: Optional[str] = None, table_type: Optional[str] = None, force: bool = False):
+        """Handle executing an up migration.
+
+        Args:
+            target_version: The version to migrate to, e.g. "v3.0.0". If not provided, the latest available version will be used.
+            table_type: The type of table to migrate. If not provided, all table types will be considered.
+        """
+
+        # If not target version is provided, use the latest available version
+        if not target_version:
+            _target_version = self.latest_schema_version
+            log_info(
+                f"No target version provided. Will migrate to the latest available version: {str(_target_version)}"
+            )
+        else:
+            _target_version = packaging_version.parse(target_version)
+
+        # Select tables to migrate
+        if table_type:
+            if table_type not in ["memory", "session", "metrics", "eval", "knowledge", "culture"]:
+                log_warning(
+                    f"Invalid table type: {table_type}. Use one of: memory, session, metrics, eval, knowledge, culture"
+                )
+                return
+            tables = [(table_type, getattr(self.db, f"{table_type}_table_name"))]
+        else:
+            tables = [
+                ("memories", self.db.memory_table_name),
+                ("sessions", self.db.session_table_name),
+                ("metrics", self.db.metrics_table_name),
+                ("evals", self.db.eval_table_name),
+                ("knowledge", self.db.knowledge_table_name),
+                ("culture", self.db.culture_table_name),
+            ]
+
+        # Handle migrations for each table separately (extend in future if needed):
+        for table_type, table_name in tables:
+            if isinstance(self.db, AsyncBaseDb):
+                current_version = packaging_version.parse(await self.db.get_latest_schema_version(table_name))
+            else:
+                current_version = packaging_version.parse(self.db.get_latest_schema_version(table_name))
+
+            if current_version is None:
+                log_info(f"Skipping migration: No version found for table {table_name}.")
+                continue
+
+            # If the target version is less or equal to the current version, no migrations needed
+            if _target_version <= current_version and not force:
+                log_info(
+                    f"Skipping migration: the version of table '{table_name}' ({current_version}) is less or equal to the target version ({_target_version})."
+                )
+                continue
+
+            log_info(
+                f"Starting database migration for table {table_name}. Current version: {current_version}. Target version: {_target_version}."
+            )
+
+            # Find files after the current version
+            latest_version = None
+            migration_executed = False
+            for version, normalised_version in self.available_versions:
+                if normalised_version > current_version:
+                    if target_version and normalised_version > _target_version:
+                        break
+
+                    log_info(f"Applying migration {normalised_version} on {table_name}")
+                    migration_executed = await self._up_migration(version, table_type, table_name)
+                    if migration_executed:
+                        log_info(f"Successfully applied migration {normalised_version} on table {table_name}")
+                    else:
+                        log_info(f"Skipping application of migration {normalised_version} on table {table_name}")
+
+                    latest_version = normalised_version.public
+
+            if migration_executed and latest_version:
+                log_info(f"Storing version {latest_version} in database for table {table_name}")
+                if isinstance(self.db, AsyncBaseDb):
+                    await self.db.upsert_schema_version(table_name, latest_version)
+                else:
+                    self.db.upsert_schema_version(table_name, latest_version)
+                log_info(f"Successfully stored version {latest_version} in database for table {table_name}")
+            log_info("----------------------------------------------------------")
+
+    async def _up_migration(self, version: str, table_type: str, table_name: str) -> bool:
+        """Run the database-specific logic to handle an up migration.
+
+        Args:
+            version: The version to migrate to, e.g. "v3.0.0"
+        """
+        migration_module = importlib.import_module(f"agno.db.migrations.versions.{version}")
+
+        try:
+            if isinstance(self.db, AsyncBaseDb):
+                return await migration_module.async_up(self.db, table_type, table_name)
+            else:
+                return migration_module.up(self.db, table_type, table_name)
+        except Exception as e:
+            log_error(f"Error running migration to version {version}: {e}")
+            raise
+
+    async def down(self, target_version: str, table_type: Optional[str] = None, force: bool = False):
+        """Handle executing a down migration.
+
+        Args:
+            target_version: The version to migrate to. e.g. "v2.3.0"
+            table_type: The type of table to migrate. If not provided, all table types will be considered.
+        """
+        _target_version = packaging_version.parse(target_version)
+
+        # Select tables to migrate
+        if table_type:
+            if table_type not in ["memory", "session", "metrics", "eval", "knowledge", "culture"]:
+                log_warning(
+                    f"Invalid table type: {table_type}. Use one of: memory, session, metrics, eval, knowledge, culture"
+                )
+                return
+            tables = [(table_type, getattr(self.db, f"{table_type}_table_name"))]
+        else:
+            tables = [
+                ("memories", self.db.memory_table_name),
+                ("sessions", self.db.session_table_name),
+                ("metrics", self.db.metrics_table_name),
+                ("evals", self.db.eval_table_name),
+                ("knowledge", self.db.knowledge_table_name),
+                ("culture", self.db.culture_table_name),
+            ]
+
+        for table_type, table_name in tables:
+            if isinstance(self.db, AsyncBaseDb):
+                current_version = packaging_version.parse(await self.db.get_latest_schema_version(table_name))
+            else:
+                current_version = packaging_version.parse(self.db.get_latest_schema_version(table_name))
+
+            if _target_version >= current_version and not force:
+                log_warning(
+                    f"Skipping down migration: the version of table '{table_name}' ({current_version}) is less or equal to the target version ({_target_version})."
+                )
+                continue
+
+            migration_executed = False
+            # Run down migration for all versions between target and current (include down of current version)
+            # Apply down migrations in reverse order to ensure dependencies are met
+            for version, normalised_version in reversed(self.available_versions):
+                if normalised_version > _target_version:
+                    log_info(f"Reverting migration {normalised_version} on table {table_name}")
+                    migration_executed = await self._down_migration(version, table_type, table_name)
+                    if migration_executed:
+                        log_info(f"Successfully reverted migration {normalised_version} on table {table_name}")
+                    else:
+                        log_info(f"Skipping revert of migration {normalised_version} on table {table_name}")
+
+            if migration_executed:
+                log_info(f"Storing version {_target_version} in database for table {table_name}")
+                if isinstance(self.db, AsyncBaseDb):
+                    await self.db.upsert_schema_version(table_name, _target_version.public)
+                else:
+                    self.db.upsert_schema_version(table_name, _target_version.public)
+                log_info(f"Successfully stored version {_target_version} in database for table {table_name}")
+
+    async def _down_migration(self, version: str, table_type: str, table_name: str) -> bool:
+        """Run the database-specific logic to handle a down migration.
+
+        Args:
+            version: The version to migrate to, e.g. "v3.0.0"
+        """
+        migration_module = importlib.import_module(f"agno.db.migrations.versions.{version}")
+        try:
+            if isinstance(self.db, AsyncBaseDb):
+                return await migration_module.async_down(self.db, table_type, table_name)
+            else:
+                return migration_module.down(self.db, table_type, table_name)
+        except Exception as e:
+            log_error(f"Error running migration to version {version}: {e}")
+            raise

agno/db/migrations/utils.py

@@ -0,0 +1,19 @@
+def quote_db_identifier(db_type: str, identifier: str) -> str:
+    """Add the right quotes to the given identifier string (table name, schema name) based on db type.
+
+    Args:
+        db_type: The database type name (e.g., "PostgresDb", "MySQLDb", "SqliteDb")
+        identifier: The identifier string to add quotes to
+
+    Returns:
+        The properly quoted identifier string
+    """
+    if db_type in ("PostgresDb", "AsyncPostgresDb"):
+        return f'"{identifier}"'
+    elif db_type in ("MySQLDb", "AsyncMySQLDb", "SingleStoreDb"):
+        return f"`{identifier}`"
+    elif db_type in ("SqliteDb", "AsyncSqliteDb"):
+        return f'"{identifier}"'
+    else:
+        # Default to double quotes for unknown types
+        return f'"{identifier}"'

agno/db/migrations/v1_to_v2.py

@@ -7,9 +7,11 @@ from typing import Any, Dict, List, Optional, Union, cast
 from sqlalchemy import text
 
 from agno.db.base import BaseDb
+from agno.db.migrations.utils import quote_db_identifier
 from agno.db.schemas.memory import UserMemory
 from agno.session import AgentSession, TeamSession, WorkflowSession
 from agno.utils.log import log_error, log_info, log_warning
+from agno.utils.string import sanitize_postgres_string, sanitize_postgres_strings
 
 
 def convert_v1_metrics_to_v2(metrics_dict: Dict[str, Any]) -> Dict[str, Any]:
@@ -45,7 +47,10 @@ def convert_v1_metrics_to_v2(metrics_dict: Dict[str, Any]) -> Dict[str, Any]:
 
 
 def convert_any_metrics_in_data(data: Any) -> Any:
-    """Recursively find and convert any metrics dictionaries and handle v1 to v2 field conversion."""
+    """Recursively find and convert any metrics dictionaries and handle v1 to v2 field conversion.
+
+    Also sanitizes all string values to remove null bytes for PostgreSQL compatibility.
+    """
     if isinstance(data, dict):
         # First apply v1 to v2 field conversion (handles extra_data extraction, thinking/reasoning_content consolidation, etc.)
         data = convert_v1_fields_to_v2(data)
@@ -62,13 +67,19 @@ def convert_any_metrics_in_data(data: Any) -> Any:
                 converted_dict[key] = convert_v1_metrics_to_v2(value)
             else:
                 converted_dict[key] = convert_any_metrics_in_data(value)
-        return converted_dict
+
+        # Sanitize all strings in the converted dict to remove null bytes
+        return sanitize_postgres_strings(converted_dict)
 
     elif isinstance(data, list):
         return [convert_any_metrics_in_data(item) for item in data]
 
+    elif isinstance(data, str):
+        # Sanitize string values to remove null bytes
+        return sanitize_postgres_string(data)
+
     else:
-        # Not a dict or list, return as-is
+        # Not a dict, list, or string, return as-is
         return data
 
 
@@ -485,15 +496,21 @@ def get_table_content_in_batches(db: BaseDb, db_schema: str, table_name: str, ba
             else:
                 raise ValueError(f"Invalid database type: {type(db).__name__}")
 
+            db_type = type(db).__name__
+            quoted_schema = (
+                quote_db_identifier(db_type=db_type, identifier=db_schema) if db_schema and db_schema.strip() else None
+            )
+            quoted_table = quote_db_identifier(db_type=db_type, identifier=table_name)
+
             offset = 0
             while True:
                 # Create a new session for each batch to avoid transaction conflicts
                 with db.Session() as sess:
                     # Handle empty schema by omitting the schema prefix (needed for SQLite)
-                    if db_schema and db_schema.strip():
-                        sql_query = f"SELECT * FROM {db_schema}.{table_name} LIMIT {batch_size} OFFSET {offset}"
+                    if quoted_schema:
+                        sql_query = f"SELECT * FROM {quoted_schema}.{quoted_table} LIMIT {batch_size} OFFSET {offset}"
                     else:
-                        sql_query = f"SELECT * FROM {table_name} LIMIT {batch_size} OFFSET {offset}"
+                        sql_query = f"SELECT * FROM {quoted_table} LIMIT {batch_size} OFFSET {offset}"
 
                     result = sess.execute(text(sql_query))
                     batch = [row._asdict() for row in result]
@@ -530,13 +547,16 @@ def get_all_table_content(db, db_schema: str, table_name: str) -> list[dict[str,
 
 
 def parse_agent_sessions(v1_content: List[Dict[str, Any]]) -> List[AgentSession]:
-    """Parse v1 Agent sessions into v2 Agent sessions and Memories"""
+    """Parse v1 Agent sessions into v2 Agent sessions and Memories
+
+    Sanitizes all string and JSON fields to remove null bytes for PostgreSQL compatibility.
+    """
     sessions_v2 = []
 
     for item in v1_content:
         session = {
             "agent_id": item.get("agent_id"),
-            "agent_data": item.get("agent_data"),
+            "agent_data": sanitize_postgres_strings(item.get("agent_data")) if item.get("agent_data") else None,
             "session_id": item.get("session_id"),
             "user_id": item.get("user_id"),
             "session_data": convert_session_data_comprehensively(item.get("session_data")),
@@ -545,6 +565,7 @@ def parse_agent_sessions(v1_content: List[Dict[str, Any]]) -> List[AgentSession]
             "created_at": item.get("created_at"),
             "updated_at": item.get("updated_at"),
         }
+        # Summary field sanitization is handled in convert_session_data_comprehensively
 
         try:
             agent_session = AgentSession.from_dict(session)
@@ -559,13 +580,16 @@ def parse_agent_sessions(v1_content: List[Dict[str, Any]]) -> List[AgentSession]
 
 
 def parse_team_sessions(v1_content: List[Dict[str, Any]]) -> List[TeamSession]:
-    """Parse v1 Team sessions into v2 Team sessions and Memories"""
+    """Parse v1 Team sessions into v2 Team sessions and Memories
+
+    Sanitizes all string and JSON fields to remove null bytes for PostgreSQL compatibility.
+    """
     sessions_v2 = []
 
     for item in v1_content:
         session = {
             "team_id": item.get("team_id"),
-            "team_data": item.get("team_data"),
+            "team_data": sanitize_postgres_strings(item.get("team_data")) if item.get("team_data") else None,
             "session_id": item.get("session_id"),
             "user_id": item.get("user_id"),
             "session_data": convert_session_data_comprehensively(item.get("session_data")),
@@ -574,6 +598,8 @@ def parse_team_sessions(v1_content: List[Dict[str, Any]]) -> List[TeamSession]:
             "created_at": item.get("created_at"),
             "updated_at": item.get("updated_at"),
         }
+        # Summary field sanitization is handled in convert_session_data_comprehensively
+
         try:
             team_session = TeamSession.from_dict(session)
         except Exception as e:
@@ -587,13 +613,18 @@ def parse_team_sessions(v1_content: List[Dict[str, Any]]) -> List[TeamSession]:
 
 
 def parse_workflow_sessions(v1_content: List[Dict[str, Any]]) -> List[WorkflowSession]:
-    """Parse v1 Workflow sessions into v2 Workflow sessions"""
+    """Parse v1 Workflow sessions into v2 Workflow sessions
+
+    Sanitizes all string and JSON fields to remove null bytes for PostgreSQL compatibility.
+    """
     sessions_v2 = []
 
     for item in v1_content:
         session = {
             "workflow_id": item.get("workflow_id"),
-            "workflow_data": item.get("workflow_data"),
+            "workflow_data": sanitize_postgres_strings(item.get("workflow_data"))
+            if item.get("workflow_data")
+            else None,
             "session_id": item.get("session_id"),
             "user_id": item.get("user_id"),
             "session_data": convert_session_data_comprehensively(item.get("session_data")),
@@ -601,9 +632,11 @@ def parse_workflow_sessions(v1_content: List[Dict[str, Any]]) -> List[WorkflowSe
             "created_at": item.get("created_at"),
             "updated_at": item.get("updated_at"),
             # Workflow v2 specific fields
-            "workflow_name": item.get("workflow_name"),
+            "workflow_name": sanitize_postgres_string(item.get("workflow_name")) if item.get("workflow_name") else None,
             "runs": convert_any_metrics_in_data(item.get("runs")),
         }
+        # Summary field sanitization is handled in convert_session_data_comprehensively
+
         try:
             workflow_session = WorkflowSession.from_dict(session)
         except Exception as e:
@@ -617,18 +650,23 @@ def parse_workflow_sessions(v1_content: List[Dict[str, Any]]) -> List[WorkflowSe
 
 
 def parse_memories(v1_content: List[Dict[str, Any]]) -> List[UserMemory]:
-    """Parse v1 Memories into v2 Memories"""
+    """Parse v1 Memories into v2 Memories
+
+    Sanitizes all string fields to remove null bytes for PostgreSQL compatibility.
+    """
     memories_v2 = []
 
     for item in v1_content:
         memory = {
             "memory_id": item.get("memory_id"),
-            "memory": item.get("memory"),
-            "input": item.get("input"),
+            "memory": sanitize_postgres_strings(item.get("memory")) if item.get("memory") else None,
+            "input": sanitize_postgres_string(item.get("input")) if item.get("input") else None,
             "updated_at": item.get("updated_at"),
             "agent_id": item.get("agent_id"),
             "team_id": item.get("team_id"),
             "user_id": item.get("user_id"),
+            "topics": sanitize_postgres_strings(item.get("topics")) if item.get("topics") else None,
+            "feedback": sanitize_postgres_string(item.get("feedback")) if item.get("feedback") else None,
         }
         memories_v2.append(UserMemory.from_dict(memory))