Merge branch 'main' into AGE-691/implement-onboarding-pages

agenta-backend/agenta_backend/migrations/postgres/README.md

@@ -19,7 +19,7 @@ docker exec -it agenta-backend-1 bash
 The above command will run the backend container on bash.
 
 ```bash
-alembic -c alembic.oss.ini revision --autogenerate -m "migration message"
+docker exec -w /app/agenta_backend/migrations/postgres agenta-backend-1 alembic -c alembic.oss.ini revision --autogenerate -m "migration message"
 ```
 
 The above command will create a script that contains the changes that was made to the database schema. Kindly update "migration message" with a message that is clear to indicate what change was made. Here are some examples:
@@ -34,4 +34,4 @@ The above command will create a script that contains the changes that was made t
 docker exec -w /app/agenta_backend/migrations/postgres agenta-backend-1 alembic -c alembic.oss.ini upgrade head
 ```
 
-The above command will be used to apply the changes in the script created to the database table(s).
+The above command will be used to apply the changes in the script created to the database table(s).

agenta-backend/agenta_backend/migrations/postgres/data_migrations/applications.py

@@ -0,0 +1,73 @@
+import os
+import uuid
+import traceback
+from typing import Optional
+
+
+import click
+from sqlalchemy.future import select
+from sqlalchemy import create_engine, delete
+from sqlalchemy.orm import sessionmaker, Session
+
+from agenta_backend.models.deprecated_models import (
+    DeprecatedEvaluatorConfigDB,
+    DeprecatedAppDB,
+)
+
+
+BATCH_SIZE = 1000
+
+
+def get_app_db(session: Session, app_id: str) -> Optional[DeprecatedAppDB]:
+    query = session.execute(select(DeprecatedAppDB).filter_by(id=uuid.UUID(app_id)))
+    return query.scalars().first()
+
+
+def update_evaluators_with_app_name():
+    engine = create_engine(os.getenv("POSTGRES_URI"))
+    sync_session = sessionmaker(engine, expire_on_commit=False)
+
+    with sync_session() as session:
+        try:
+            offset = 0
+            while True:
+                records = (
+                    session.execute(
+                        select(DeprecatedEvaluatorConfigDB)
+                        .filter(DeprecatedEvaluatorConfigDB.app_id.isnot(None))
+                        .offset(offset)
+                        .limit(BATCH_SIZE)
+                    )
+                    .scalars()
+                    .all()
+                )
+                if not records:
+                    break
+
+                # Update records with app_name as prefix
+                for record in records:
+                    evaluator_config_app = get_app_db(
+                        session=session, app_id=str(record.app_id)
+                    )
+                    if record.app_id is not None and evaluator_config_app is not None:
+                        record.name = f"{record.name} ({evaluator_config_app.app_name})"
+
+                session.commit()
+                offset += BATCH_SIZE
+
+            # Delete deprecated evaluator configs with app_id as None
+            session.execute(
+                delete(DeprecatedEvaluatorConfigDB).where(
+                    DeprecatedEvaluatorConfigDB.app_id.is_(None)
+                )
+            )
+            session.commit()
+        except Exception as e:
+            session.rollback()
+            click.echo(
+                click.style(
+                    f"ERROR updating evaluator config names: {traceback.format_exc()}",
+                    fg="red",
+                )
+            )
+            raise e

agenta-backend/agenta_backend/migrations/postgres/data_migrations/projects.py

@@ -0,0 +1,188 @@
+import os
+import traceback
+from typing import Sequence
+
+
+import click
+from sqlalchemy.future import select
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, Session
+
+from agenta_backend.models.db_models import (
+    ProjectDB,
+    AppDB,
+    AppVariantDB,
+    AppVariantRevisionsDB,
+    VariantBaseDB,
+    DeploymentDB,
+    ImageDB,
+    AppEnvironmentDB,
+    AppEnvironmentRevisionDB,
+    EvaluationScenarioDB,
+    EvaluationDB,
+    EvaluatorConfigDB,
+    HumanEvaluationDB,
+    HumanEvaluationScenarioDB,
+    TestSetDB,
+)
+
+
+BATCH_SIZE = 1000
+MODELS = [
+    AppDB,
+    AppVariantDB,
+    AppVariantRevisionsDB,
+    VariantBaseDB,
+    DeploymentDB,
+    ImageDB,
+    AppEnvironmentDB,
+    AppEnvironmentRevisionDB,
+    EvaluationScenarioDB,
+    EvaluationDB,
+    EvaluatorConfigDB,
+    HumanEvaluationDB,
+    HumanEvaluationScenarioDB,
+    TestSetDB,
+]
+
+
+def get_default_projects(session):
+    query = session.execute(select(ProjectDB).filter_by(is_default=True))
+    return query.scalars().all()
+
+
+def check_for_multiple_default_projects(session: Session) -> Sequence[ProjectDB]:
+    default_projects = get_default_projects(session)
+    if len(default_projects) > 1:
+        raise ValueError(
+            "Multiple default projects found. Please ensure only one exists."
+        )
+    return default_projects
+
+
+def create_default_project():
+    PROJECT_NAME = "Default Project"
+    engine = create_engine(os.getenv("POSTGRES_URI"))
+    sync_session = sessionmaker(engine, expire_on_commit=False)
+
+    with sync_session() as session:
+        try:
+            default_projects = check_for_multiple_default_projects(session)
+            if len(default_projects) == 0:
+                new_project = ProjectDB(project_name=PROJECT_NAME, is_default=True)
+                session.add(new_project)
+                session.commit()
+
+        except Exception as e:
+            session.rollback()
+            click.echo(
+                click.style(
+                    f"ERROR creating default project: {traceback.format_exc()}",
+                    fg="red",
+                )
+            )
+            raise e
+
+
+def remove_default_project():
+    engine = create_engine(os.getenv("POSTGRES_URI"))
+    sync_session = sessionmaker(engine, expire_on_commit=False)
+
+    with sync_session() as session:
+        try:
+            default_projects = check_for_multiple_default_projects(session)
+            if len(default_projects) == 0:
+                click.echo(
+                    click.style("No default project found to remove.", fg="yellow")
+                )
+                return
+
+            session.delete(default_projects[0])
+            session.commit()
+            click.echo(click.style("Default project removed successfully.", fg="green"))
+
+        except Exception as e:
+            session.rollback()
+            click.echo(click.style(f"ERROR: {traceback.format_exc()}", fg="red"))
+            raise e
+
+
+def add_project_id_to_db_entities():
+    engine = create_engine(os.getenv("POSTGRES_URI"))
+    sync_session = sessionmaker(engine, expire_on_commit=False)
+
+    with sync_session() as session:
+        try:
+            default_project = check_for_multiple_default_projects(session)[0]
+            for model in MODELS:
+                offset = 0
+                while True:
+                    records = (
+                        session.execute(
+                            select(model)
+                            .where(model.project_id == None)
+                            .offset(offset)
+                            .limit(BATCH_SIZE)
+                        )
+                        .scalars()
+                        .all()
+                    )
+                    if not records:
+                        break
+
+                    # Update records with default project_id
+                    for record in records:
+                        record.project_id = default_project.id
+
+                    session.commit()
+                    offset += BATCH_SIZE
+
+        except Exception as e:
+            session.rollback()
+            click.echo(
+                click.style(
+                    f"ERROR adding project_id to db entities: {traceback.format_exc()}",
+                    fg="red",
+                )
+            )
+            raise e
+
+
+def remove_project_id_from_db_entities():
+    engine = create_engine(os.getenv("POSTGRES_URI"))
+    sync_session = sessionmaker(engine, expire_on_commit=False)
+
+    with sync_session() as session:
+        try:
+            for model in MODELS:
+                offset = 0
+                while True:
+                    records = (
+                        session.execute(
+                            select(model)
+                            .where(model.project_id != None)
+                            .offset(offset)
+                            .limit(BATCH_SIZE)
+                        )
+                        .scalars()
+                        .all()
+                    )
+                    if not records:
+                        break
+
+                    # Update records project_id column with None
+                    for record in records:
+                        record.project_id = None
+
+                    session.commit()
+                    offset += BATCH_SIZE
+
+        except Exception as e:
+            session.rollback()
+            click.echo(
+                click.style(
+                    f"ERROR removing project_id to db entities: {traceback.format_exc()}",
+                    fg="red",
+                )
+            )
+            raise e

agenta-backend/agenta_backend/migrations/postgres/env.py

@@ -47,6 +47,7 @@ def run_migrations_offline() -> None:
     url = config.get_main_option("sqlalchemy.url")
     context.configure(
         url=url,
+        transaction_per_migration=True,
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
@@ -57,7 +58,11 @@ def run_migrations_offline() -> None:
 
 
 def do_run_migrations(connection: Connection) -> None:
-    context.configure(connection=connection, target_metadata=target_metadata)
+    context.configure(
+        transaction_per_migration=True,
+        connection=connection,
+        target_metadata=target_metadata,
+    )
 
     with context.begin_transaction():
         context.run_migrations()

agenta-backend/agenta_backend/migrations/postgres/utils.py

@@ -5,13 +5,14 @@
 
 import click
 import asyncpg
+
+from sqlalchemy import inspect, text, Engine
 from sqlalchemy.exc import ProgrammingError
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
 
 from alembic import command
 from alembic.config import Config
-from sqlalchemy import inspect, text
 from alembic.script import ScriptDirectory
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
 
 from agenta_backend.utils.common import isCloudEE, isCloudDev
 
@@ -173,3 +174,31 @@ async def check_if_templates_table_exist():
             await engine.dispose()
 
         return True
+
+
+def unique_constraint_exists(
+    engine: Engine, table_name: str, constraint_name: str
+) -> bool:
+    """
+    The function checks if a unique constraint with a specific name exists on a table in a PostgreSQL
+    database.
+
+    Args:
+        - engine (Engine): instance of a database engine that represents a connection to a database.
+        - table_name (str): name of the table to check the existence of the unique constraint.
+        - constraint_name (str): name of the unique constraint to check for existence.
+
+    Returns:
+        - returns a boolean value indicating whether a unique constraint with the specified `constraint_name` exists in the table.
+    """
+
+    with engine.connect() as conn:
+        result = conn.execute(
+            text(
+                f"""
+        SELECT conname FROM pg_constraint
+        WHERE conname = '{constraint_name}' AND conrelid = '{table_name}'::regclass;
+        """
+            )
+        )
+        return result.fetchone() is not None

agenta-backend/agenta_backend/migrations/postgres/versions/22d29365f5fc_update_evaluators_names_with_app_name_.py

@@ -0,0 +1,32 @@
+"""Update evaluators names with app name as prefix
+
+Revision ID: 22d29365f5fc
+Revises: 6cfe239894fb
+Create Date: 2024-09-16 11:38:33.886908
+
+"""
+
+from typing import Sequence, Union
+
+from agenta_backend.migrations.postgres.data_migrations.applications import (
+    update_evaluators_with_app_name,
+)
+
+
+# revision identifiers, used by Alembic.
+revision: str = "22d29365f5fc"
+down_revision: Union[str, None] = "6cfe239894fb"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### custom command ###
+    update_evaluators_with_app_name()
+    # ### end custom command ###
+
+
+def downgrade() -> None:
+    # ### custom command ###
+    pass
+    # ### end custom command ###