feat: auto-downgrading revision

.github/workflows/test-pr-e2e.yml

@@ -300,7 +300,7 @@ jobs:
       is-fork: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork }}
       backend-image-name: ${{ needs.build-backend.outputs.image_name }}
       frontend-image-name: ${{ needs.build-frontend.outputs.image_name }}
-  
+
   run-postgresql-without-redis:
     needs: [build-frontend, build-backend, prepare-test-environment]
     uses: ./.github/workflows/run-e2e-tests.yml
@@ -311,7 +311,7 @@ jobs:
       is-fork: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork }}
       backend-image-name: ${{ needs.build-backend.outputs.image_name }}
       frontend-image-name: ${{ needs.build-frontend.outputs.image_name }}
-  
+
   run-sqlite-without-redis:
     needs: [build-frontend, build-backend, prepare-test-environment]
     uses: ./.github/workflows/run-e2e-tests.yml
@@ -321,4 +321,15 @@ jobs:
       python-version: 3.11
       is-fork: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork }}
       backend-image-name: ${{ needs.build-backend.outputs.image_name }}
+      frontend-image-name: ${{ needs.build-frontend.outputs.image_name }}
+
+  run-migrations-tests:
+    needs: [build-frontend, build-backend, prepare-test-environment]
+    uses: ./.github/workflows/run-migrations-e2e-tests.yml
+    with:
+      db-type: postgres-migrations
+      redis_enabled: false
+      python-version: 3.11
+      is-fork: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork }}
+      backend-image-name: ${{ needs.build-backend.outputs.image_name }}
       frontend-image-name: ${{ needs.build-frontend.outputs.image_name }}

docker-compose-with-arq.yml

@@ -24,6 +24,7 @@ services:
       - REDIS_PORT=6379
     volumes:
       - ./state:/state
+      - ./tmp/keep/migrations:/tmp/keep/migrations
     depends_on:
       - keep-arq-redis
 

docker-compose-with-auth.yml

@@ -25,6 +25,7 @@ services:
       - KEEP_DEFAULT_PASSWORD=keep
     volumes:
       - ./state:/state
+      - ./tmp/keep/migrations:/tmp/keep/migrations
 
   keep-websocket-server:
     extends:

docker-compose-with-otel.yaml

@@ -111,6 +111,7 @@ services:
     volumes:
       - .:/app
       - ./state:/state
+      - ./tmp/keep/migrations:/tmp/keep/migrations
 
   keep-websocket-server:
     extends:

docker-compose.common.yml

@@ -20,8 +20,10 @@ services:
       - PORT=8080
       - SECRET_MANAGER_TYPE=FILE
       - SECRET_MANAGER_DIRECTORY=/state
+      - MIGRATIONS_PATH=/tmp/keep/migrations
       - DATABASE_CONNECTION_STRING=sqlite:////state/db.sqlite3?check_same_thread=False
       - OPENAI_API_KEY=$OPENAI_API_KEY
+      - ALLOW_DB_DOWNGRADE=false
       - PUSHER_APP_ID=1
       - PUSHER_APP_KEY=keepappkey
       - PUSHER_APP_SECRET=keepappsecret

docker-compose.dev.yml

@@ -24,6 +24,7 @@ services:
     volumes:
       - .:/app
       - ./state:/state
+      - ./tmp/keep/migrations:/tmp/keep/migrations
 
   keep-websocket-server:
     extends:

docker-compose.yml

@@ -23,6 +23,7 @@ services:
       - KEEP_METRICS=true
     volumes:
       - ./state:/state
+      - ./tmp/keep/migrations:/tmp/keep/migrations
 
   keep-websocket-server:
     extends:

docs/deployment/configuration.mdx

@@ -28,6 +28,7 @@ Keep is highly configurable through environment variables. This allows you to cu
 |      **KEEP_STORE_RAW_ALERTS**       |             Enables storing of raw alerts             |    No    |            "false"             |      "true" or "false"       |
 | **TENANT_CONFIGURATION_RELOAD_TIME** |    Time in minutes to reload tenant configurations    |    No    |               5                |       Positive integer       |
 |       **KEEP_LIVE_DEMO_MODE**        | Keep will simulate incoming alerts and other activity |    No    |            "false"             |      "true" or "false"       |
+|         **MIGRATIONS_PATH**          |             Path to migrations directory              |    No    |     "/tmp/keep/migrations"     |        Valid os path         |
 
 ### Logging and Environment
 
@@ -63,6 +64,7 @@ Keep is highly configurable through environment variables. This allows you to cu
 |     **DB_SERVICE_ACCOUNT**     |    Service account for database impersonation     |    No    |               None                |    Valid service account email     |
 |         **DB_IP_TYPE**         |          Specifies the Cloud SQL IP type          |    No    |             "public"              |    "public", "private" or "psc"    |
 |      **SKIP_DB_CREATION**      |      Skips database creation and migrations       |    No    |              "false"              |         "true" or "false"          |
+|     **ALLOW_DB_DOWNGRADE**     |       Enables downgrading database schema         |    No    |              "false"              |         "true" or "false"          |
 
 ### Resource Provisioning
 

keep/api/core/db_on_start.py

@@ -16,9 +16,11 @@
 import hashlib
 import logging
 import os
+import shutil
 
 import alembic.command
 import alembic.config
+from alembic.runtime.migration import MigrationContext
 from sqlalchemy.exc import IntegrityError
 from sqlmodel import Session, select
 
@@ -167,23 +169,123 @@ def try_create_single_tenant(tenant_id: str, create_default_user=True) -> None:
             logger.exception("Failed to create single tenant")
             pass
 
+def get_current_revision():
+    """Get current app revision"""
+    with engine.connect() as connection:
+        context = MigrationContext.configure(connection)
+        return context.get_current_revision()
 
-def migrate_db():
+def copy_migrations(app_migrations_path, local_migrations_path):
+    """Copy migrations to a local backup folder for safe downgrade purposes."""
+
+    source_versions_path = os.path.join(app_migrations_path, "versions")
+
+    # Ensure destination exists
+    try:
+        os.makedirs(local_migrations_path, exist_ok=True)
+    except Exception as e:
+        logger.error(f"Failed to create local migrations folder with error: {e}")
+
+
+    # Clear previous versioned migrations to ensure only migrations relevant to the current version are present
+    for filename in os.listdir(local_migrations_path):
+        file_path = os.path.join(local_migrations_path, filename)
+        if os.path.isfile(file_path) or os.path.islink(file_path):
+            os.remove(file_path)
+
+    # Alembic needs the full migration history to safely perform a downgrade to earlier versions
+    # Copy new migrations
+    for item in os.listdir(source_versions_path):
+        src = os.path.join(source_versions_path, item)
+        dst = os.path.join(local_migrations_path, item)
+        if os.path.isdir(src):
+            shutil.copytree(src, dst, dirs_exist_ok=True)
+        else:
+            shutil.copy(src, dst)
+
+def downgrade_db(config, expected_revision, local_migrations_path, app_migrations_path):
+    """
+    Downgrade the DB to the previous revision, using local backup migrations temporarily.
+    Restores original migrations after downgrade.
+    """
+    source_versions_path = os.path.join(app_migrations_path, "versions")
+    source_versions_path_copy = os.path.join(app_migrations_path, "versions_copy")
+
+    try:
+        logger.info("Backing up original migrations...")
+        if os.path.exists(source_versions_path_copy):
+            shutil.rmtree(source_versions_path_copy)
+        shutil.move(source_versions_path, source_versions_path_copy)
+        logger.info("Original migrations backed up.")
+
+        logger.info("Restoring migrations from local backup...")
+        shutil.copytree(local_migrations_path, source_versions_path)
+        logger.info("Migrations restored from local.")
+
+        logger.info("Downgrading the database...")
+        alembic.command.downgrade(config, expected_revision)
+        logger.info("Database successfully downgraded.")
+
+    except Exception as e:
+        logger.error(f"Error occurred during downgrade process: {e}")
+    finally:
+        logger.info("Restoring original migrations...")
+        try:
+            if os.path.exists(source_versions_path):
+                shutil.rmtree(source_versions_path)
+            if os.path.exists(source_versions_path_copy):
+                shutil.move(source_versions_path_copy, source_versions_path)
+                logger.info("Original migrations restored!")
+            else:
+                logger.warning("Backup not found!!! Original migrations not restored!!!")
+        except Exception as restore_error:
+            logger.error(f"Failed to restore original migrations: {restore_error}")
+
+def migrate_db(config_path: str = None, app_migrations_path: str = None):
     """
     Run migrations to make sure the DB is up-to-date.
     """
     if os.environ.get("SKIP_DB_CREATION", "false") == "true":
         logger.info("Skipping running migrations...")
         return None
 
-    logger.info("Running migrations...")
-    config_path = os.path.dirname(os.path.abspath(__file__)) + "/../../" + "alembic.ini"
+    config_path = config_path or os.path.dirname(os.path.abspath(__file__)) + "/../../" + "alembic.ini"
     config = alembic.config.Config(file_=config_path)
     # Re-defined because alembic.ini uses relative paths which doesn't work
     # when running the app as a pyhton pakage (could happen form any path)
+
+    # This path will be used to save migrations locally for safe downgrade purposes
+    local_migrations_path = os.environ.get("MIGRATIONS_PATH", "/tmp/keep/migrations")
+    app_migrations_path = app_migrations_path or os.path.dirname(os.path.abspath(__file__)) + "/../models/db/migrations"
     config.set_main_option(
         "script_location",
-        os.path.dirname(os.path.abspath(__file__)) + "/../models/db/migrations",
+        app_migrations_path,
     )
-    alembic.command.upgrade(config, "head")
+    alembic_script = alembic.script.ScriptDirectory.from_config(config)
+
+    current_revision = get_current_revision()
+    expected_revision = alembic_script.get_current_head()
+
+    # If the current revision is the same as the expected revision, we don't need to run migrations
+    if current_revision and expected_revision and current_revision == expected_revision:
+        logger.info("Database schema is up-to-date!")
+        return None
+
+    logger.warning(f"Database schema ({current_revision}) doesn't match application version ({expected_revision})")
+    logger.info("Running migrations...")
+    try:
+        alembic.command.upgrade(config, "head")
+    except Exception as e:
+        logger.error(f"{e} it's seems like Keep was rolled back to a previous version")
+
+        if not os.getenv("ALLOW_DB_DOWNGRADE", "false") == "true":
+            logger.error(f"ALLOW_DB_DOWNGRADE is not set to true, but the database schema ({current_revision}) doesn't match application version ({expected_revision})")
+            raise RuntimeError("Database downgrade is not allowed")
+
+        logger.info("Downgrading database schema...")
+        downgrade_db(config, expected_revision, local_migrations_path, app_migrations_path)
+
+    # Copy migrations to local folder for safe downgrade purposes
+    copy_migrations(app_migrations_path, local_migrations_path)
+
     logger.info("Finished migrations")

pyproject.toml

@@ -86,6 +86,7 @@ psycopg-binary = "^3.2.3"
 psycopg = "^3.2.3"
 prometheus-client = "^0.21.1"
 psycopg2-binary = "^2.9.10"
+black = "^24.3.0"
 
 prometheus-fastapi-instrumentator = "^7.0.0"
 slowapi = "^0.1.9"
@@ -100,7 +101,6 @@ awscli = "^1.40.8"
 pre-commit = "^3.0.4"
 pre-commit-hooks = "^4.4.0"
 yamllint = "^1.29.0"
-black = "^24.3.0"
 isort = "^5.12.0"
 autopep8 = "^2.0.1"
 flake8 = "^6.0.0"

tests/e2e_tests/docker-compose-e2e-mysql.yml

@@ -29,21 +29,23 @@ services:
       - API_URL=http://keep-backend:8080
       - POSTHOG_DISABLED=true
       - SENTRY_DISABLED=true
+      - ALLOW_DB_DOWNGRADE=true
 
   # Backend Services
   keep-backend:
     # to be replaced in github actions
     image: "%KEEPBACKEND_IMAGE%"
+    ports:
+      - "8080:8080"
     environment:
       - AUTH_TYPE=NO_AUTH
       - DATABASE_CONNECTION_STRING=mysql+pymysql://root:keep@keep-database:3306/keep
       - POSTHOG_DISABLED=true
       - SECRET_MANAGER_DIRECTORY=/app
+      - MIGRATIONS_PATH=/tmp/migrations
       - SQLALCHEMY_WARN_20=1
       - REDIS=${REDIS:-false}
       - REDIS_HOST=${REDIS_HOST:-localhost}
-    ports:
-      - "8080:8080"
     depends_on:
       keep-database:
         condition: service_healthy