Merge pull request #1966 from Agenta-AI/feature/apply-migrations-for-…

…first-time-users

[Enhancement] Enable automatic migration for first-time users and optionally for returning users

.github/workflows/run-backend-tests.yml

@@ -33,9 +33,6 @@ jobs:
         run: |
           sleep 10 && curl -i http://localhost/api/openapi.json
 
-      - name: Apply Schema Migration
-        run:  docker exec -w /app/agenta_backend/migrations/postgres agenta-backend-test alembic -c alembic.oss.ini upgrade head
-
       - name: Restart Backend Service To Fetch Template Images
         run: docker container restart agenta-backend-test
 

.github/workflows/run-frontend-tests.yml

@@ -46,9 +46,6 @@ jobs:
           curl http://localhost/api/organizations/own/
           curl http://localhost/api/containers/templates/
 
-      - name: Apply Schema Migration
-        run:  docker exec -w /app/agenta_backend/migrations/postgres agenta-backend-test alembic -c alembic.oss.ini upgrade head
-
       - name: Set Node.js 18
         uses: actions/setup-node@v4
         with:

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

@@ -31,7 +31,7 @@ The above command will create a script that contains the changes that was made t
 ### Applying Migrations
 
 ```bash
-alembic -c alembic.oss.ini upgrade head
+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).

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

@@ -1,10 +1,13 @@
 import os
+import asyncio
 import logging
+import traceback
 
 import click
 import asyncpg
 from sqlalchemy.exc import ProgrammingError
 
+from alembic import command
 from alembic.config import Config
 from sqlalchemy import inspect, text
 from alembic.script import ScriptDirectory
@@ -49,11 +52,6 @@ def is_initial_setup(engine) -> bool:
     # Check if all required tables exist in the database
     all_tables_exist = all(table in existing_tables for table in required_tables)
 
-    # Log the status of the tables
-    logger.info(f"Required tables: {required_tables}")
-    logger.info(f"Existing tables: {existing_tables}")
-    logger.info(f"All tables exist: {all_tables_exist}")
-
     return not all_tables_exist
 
 
@@ -105,6 +103,44 @@ async def get_pending_migrations():
     return pending_migrations
 
 
+def run_alembic_migration():
+    """
+    Applies migration for first-time users and also checks the environment variable "AGENTA_AUTO_MIGRATIONS" to determine whether to apply migrations for returning users.
+    """
+
+    try:
+        pending_migrations = asyncio.run(get_pending_migrations())
+        APPLY_AUTO_MIGRATIONS = os.environ.get("AGENTA_AUTO_MIGRATIONS")
+        FIRST_TIME_USER = True if "alembic_version" in pending_migrations else False
+
+        if FIRST_TIME_USER or APPLY_AUTO_MIGRATIONS == "true":
+            command.upgrade(alembic_cfg, "head")
+            click.echo(
+                click.style(
+                    "\nMigration applied successfully. The container will now exit.",
+                    fg="green",
+                ),
+                color=True,
+            )
+        else:
+            click.echo(
+                click.style(
+                    "\nAll migrations are up-to-date. The container will now exit.",
+                    fg="yellow",
+                ),
+                color=True,
+            )
+    except Exception as e:
+        click.echo(
+            click.style(
+                f"\nAn ERROR occured while applying migration: {traceback.format_exc()}\nThe container will now exit.",
+                fg="red",
+            ),
+            color=True,
+        )
+        raise e
+
+
 async def check_for_new_migrations():
     """
     Checks for new migrations and notify the user.
@@ -114,7 +150,7 @@ async def check_for_new_migrations():
     if len(pending_migrations) >= 1:
         click.echo(
             click.style(
-                f"\nWe have detected that there are pending database migrations {pending_migrations} that need to be applied to keep your system up to date. \nTo ensure your application functions correctly with the latest updates, please follow the guide here => https://docs.agenta.ai/self-host/migration/applying-schema-migration\n",
+                f"\nWe have detected that there are pending database migrations {pending_migrations} that need to be applied to keep the application up to date. To ensure the application functions correctly with the latest updates, please follow the guide here => https://docs.agenta.ai/self-host/migration/applying-schema-migration\n",
                 fg="yellow",
             ),
             color=True,

docker-compose.gh.yml

@@ -60,8 +60,26 @@ services:
         depends_on:
             postgres:
                 condition: service_healthy
+            apply_alembic_migration:
+                condition: service_completed_successfully
         restart: always
 
+    apply_alembic_migration:
+        build: ghcr.io/agenta-ai/agenta-backend
+        command: sh -c "python -c 'from agenta_backend.migrations.postgres.utils import run_alembic_migration; run_alembic_migration()'"
+        environment:
+            - FEATURE_FLAG=oss
+            - POSTGRES_URI=postgresql+asyncpg://username:password@postgres:5432/agenta_oss
+            - ALEMBIC_CFG_PATH=/app/agenta_backend/migrations/postgres/alembic.oss.ini
+            - AGENTA_AUTO_MIGRATIONS=false
+        volumes:
+            - /var/run/docker.sock:/var/run/docker.sock
+        depends_on:
+          postgres:
+            condition: service_healthy
+        networks:
+            - agenta-network
+
     agenta-web:
         container_name: agenta-web-1
         image: ghcr.io/agenta-ai/agenta-web

docker-compose.prod.yml

@@ -63,8 +63,27 @@ services:
         depends_on:
             postgres:
                 condition: service_healthy
+            apply_alembic_migration:
+                condition: service_completed_successfully
         restart: always
 
+    apply_alembic_migration:
+        build: ./agenta-backend
+        command: sh -c "python -c 'from agenta_backend.migrations.postgres.utils import run_alembic_migration; run_alembic_migration()'"
+        environment:
+            - FEATURE_FLAG=oss
+            - POSTGRES_URI=postgresql+asyncpg://username:password@postgres:5432/agenta_oss
+            - ALEMBIC_CFG_PATH=/app/agenta_backend/migrations/postgres/alembic.oss.ini
+            - AGENTA_AUTO_MIGRATIONS=false
+        volumes:
+            - /var/run/docker.sock:/var/run/docker.sock
+            - ./agenta-backend/agenta_backend:/app/agenta_backend
+        depends_on:
+          postgres:
+            condition: service_healthy
+        networks:
+            - agenta-network
+
     agenta-web:
         build:
             context: ./agenta-web

docker-compose.test.yml

@@ -60,11 +60,30 @@ services:
         depends_on:
             postgres:
                 condition: service_healthy
+            apply_alembic_migration:
+                condition: service_completed_successfully
         extra_hosts:
             - host.docker.internal:host-gateway
         networks:
             - agenta-network
 
+    apply_alembic_migration:
+        build: ./agenta-backend
+        command: sh -c "python -c 'from agenta_backend.migrations.postgres.utils import run_alembic_migration; run_alembic_migration()'"
+        environment:
+            - FEATURE_FLAG=oss
+            - POSTGRES_URI=postgresql+asyncpg://username:password@postgres:5432
+            - ALEMBIC_CFG_PATH=/app/agenta_backend/migrations/postgres/alembic.oss.ini
+            - AGENTA_AUTO_MIGRATIONS=true
+        volumes:
+            - /var/run/docker.sock:/var/run/docker.sock
+            - ./agenta-backend/agenta_backend:/app/agenta_backend
+        depends_on:
+          postgres:
+            condition: service_healthy
+        networks:
+            - agenta-network
+
     agenta-web:
         build:
             context: ./agenta-web

docker-compose.yml

@@ -64,8 +64,28 @@ services:
         depends_on:
             postgres:
                 condition: service_healthy
+            apply_alembic_migration:
+                condition: service_completed_successfully
         restart: always
 
+    apply_alembic_migration:
+        build: ./agenta-backend
+        command: sh -c "python -c 'from agenta_backend.migrations.postgres.utils import run_alembic_migration; run_alembic_migration()'"
+        environment:
+            - FEATURE_FLAG=oss
+            - POSTGRES_URI=postgresql+asyncpg://username:password@postgres:5432/agenta_oss
+            - ALEMBIC_CFG_PATH=/app/agenta_backend/migrations/postgres/alembic.oss.ini
+            - AGENTA_AUTO_MIGRATIONS=false
+        volumes:
+            - /var/run/docker.sock:/var/run/docker.sock
+            - ./agenta-backend/migrations:/app/migrations
+            - ./agenta-backend/agenta_backend:/app/agenta_backend
+        depends_on:
+          postgres:
+            condition: service_healthy
+        networks:
+            - agenta-network
+
     agenta-web:
         build:
             context: ./agenta-web

docs/self-host/migration/applying-schema-migration.mdx

@@ -7,7 +7,7 @@ This guide provides step-by-step instructions for applying schema migration to y
 
 Schema migration is different from [Migration to PostgreSQL](https://docs.agenta.ai/self-host/migration/migration-to-postgres). Schema migration involves modifying the database schema (such as creating, altering, or dropping tables and columns) to match the application's current requirements.
 
-### Steps-by-Steps Guide
+### Applying the Migration
 
 To apply schema migrations, you can use the following command. This sets the working directory to `/app/agenta_backend/migrations/postgres` in the backend container before executing the Alembic command. It ensures that Alembic looks for the configuration file `alembic.oss.ini` and other necessary files in the specified directory.
 
@@ -20,3 +20,7 @@ docker exec -w /app/agenta_backend/migrations/postgres agenta-backend-1 alembic
 After completing the migration, ensure you check the data integrity in PostgreSQL by accessing Agenta on the web and verifying that your data is intact and everything works fine.
 
 In the event that you encounter issues and need to revert the schema migration, you can revert by running `alembic -c alembic.oss.ini downgrade head`. Afterwards, create a GitHub issue describing the problem you encountered.
+
+### Auto Migration
+
+In some cases, you would prefer that these migrations are applied to the application automatically without you running the `alembic upgrade` command. In this case, you need to update the value of `AGENTA_AUTO_MIGRATION` to `true` in `apply_alembic_migration` service, located in the compose file you're about to run (or are running already).