keephq · korolenkowork · May 1, 2025 · May 1, 2025 · May 1, 2025 · May 1, 2025
diff --git a/docker-compose.common.yml b/docker-compose.common.yml
@@ -22,6 +22,7 @@ services:
       - SECRET_MANAGER_DIRECTORY=/state
       - 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

diff --git a/docs/deployment/configuration.mdx b/docs/deployment/configuration.mdx
@@ -63,6 +63,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
 

diff --git a/keep/api/core/db_on_start.py b/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,88 @@ 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")
+
+    # If the destination folder already exists, remove it
+    if os.path.exists(local_migrations_path):
+        shutil.rmtree(local_migrations_path)
+
+    # Copy the entire directory tree
+    shutil.copytree(source_versions_path, local_migrations_path)
+
+def downgrade_db(config, expected_revision, local_migrations_path, app_migrations_path):
+    """
+    Downgrade the DB to the previous revision.
+    """
+    source_versions_path = os.path.join(app_migrations_path, "versions")
+    source_versions_path_copy = os.path.join(app_migrations_path, "versions_copy")
+
+    # Rename the source versions folder before restoring migrations from the backup folder
+    shutil.move(source_versions_path, source_versions_path_copy)
+
+    # Restore migrations from the backup folder
+    shutil.copytree(local_migrations_path, source_versions_path)
+
+    # Downgrade database from backup migrations
+    alembic.command.downgrade(config, expected_revision)
+
+    # Restoring source migrations
+    shutil.rmtree(source_versions_path)
+    shutil.move(source_versions_path_copy, source_versions_path)
+
+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("SECRET_MANAGER_DIRECTORY", "/state") + "/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 KeepHQ 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")
diff --git a/tests/migrations/__init__.py b/tests/migrations/__init__.py
diff --git a/tests/migrations/alembic.ini b/tests/migrations/alembic.ini
@@ -0,0 +1,51 @@
+[alembic]
+# Re-defined in the keep/api/core/db_on_start.py to make it stable while keep is installed as a package
+script_location = keep/api/models/db/migrations
+file_template = %%(year)d-%%(month).2d-%%(day).2d-%%(hour).2d-%%(minute).2d_%%(rev)s
+prepend_sys_path = .
+output_encoding = utf-8
+
+
+[post_write_hooks]
+hooks = black,isort
+
+black.type = console_scripts
+black.entrypoint = black
+
+isort.type = console_scripts
+isort.entrypoint = isort
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [PID %(process)d] [%(name)s] %(message)s
+datefmt = %H:%M:%S
diff --git a/tests/migrations/env.py b/tests/migrations/env.py
@@ -0,0 +1,136 @@
+import asyncio
+import os
+from logging.config import fileConfig
+
+from alembic import context
+from alembic.script import ScriptDirectory
+from sqlalchemy import create_engine, StaticPool
+from sqlalchemy.future import Connection
+from sqlmodel import SQLModel
+
+import keep.api.logging
+from keep.api.core.db_utils import create_db_engine
+from keep.api.models.db.action import *
+from keep.api.models.db.ai_suggestion import *
+from keep.api.models.db.alert import *
+from keep.api.models.db.dashboard import *
+from keep.api.models.db.extraction import *
+from keep.api.models.db.facet import *
+from keep.api.models.db.maintenance_window import *
+from keep.api.models.db.mapping import *
+from keep.api.models.db.preset import *
+from keep.api.models.db.provider import *
+from keep.api.models.db.rule import *
+from keep.api.models.db.statistics import *
+from keep.api.models.db.tenant import *
+from keep.api.models.db.topology import *
+from keep.api.models.db.user import *
+from keep.api.models.db.workflow import *
+
+target_metadata = SQLModel.metadata
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    # backup the current config
+    logging_config = config.get_section("loggers")
+    fileConfig(config.config_file_name)
+
+
+async def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    connectable = create_engine(
+            "sqlite:///:memory:",
+            connect_args={"check_same_thread": False},
+            poolclass=StaticPool,
+        )
+    context.configure(
+        url=str(connectable.url),
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        render_as_batch=True,
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def do_run_migrations(connection: Connection) -> None:
+    """
+    Run actual sync migrations.
+
+    :param connection: connection to the database.
+    """
+    context.configure(
+        connection=connection, target_metadata=target_metadata, render_as_batch=True
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+async def run_migrations_online() -> None:
+    """
+    Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+    """
+    connectable = create_db_engine()
+    try:
+        do_run_migrations(connectable.connect())
+    except Exception as e:
+        # print all migrations so we will know what failed
+        list_migrations(connectable)
+        raise e
+
+
+def list_migrations(connectable):
+    """
+    List all migrations and their status for debugging.
+    """
+    try:
+        # Get the script directory from the alembic context
+        script_directory = ScriptDirectory.from_config(config)
+        current_rev = script_directory.get_current_head()
+        # List all available migrations
+        pid = os.getpid()
+        print(f"[{pid}] Available migrations:")
+        try:
+            for script in script_directory.walk_revisions():
+                status = (
+                    "PENDING"
+                    if current_rev and script.revision > current_rev
+                    else "APPLIED"
+                )
+                print(f"  - {script.revision}: {script.doc} ({status})")
+        except Exception as exc:
+            logger.exception(f"Failed to list migrations: {exc}")
+    except Exception as exc:
+        logger.exception(f"Failed to process migration information: {exc}")
+
+
+loop = asyncio.get_event_loop()
+if context.is_offline_mode():
+    task = run_migrations_offline()
+else:
+    task = run_migrations_online()
+
+loop.run_until_complete(task)
+# SHAHAR: set back the logs to the default after alembic is done
+keep.api.logging.setup_logging()
diff --git a/tests/migrations/script.py.mako b/tests/migrations/script.py.mako
@@ -0,0 +1,27 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+import sqlalchemy_utils
+
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}