feat: Hello Migrations 👋 (#1368)

docs/development/getting-started.mdx

@@ -29,6 +29,15 @@ poetry run playwright install;
 poetry run coverage run --branch -m pytest -s tests/e2e_tests/
 ```
 
+### Migrations
+
+Migrations are automatically executed on a server startup. To create a migration:
+```bash
+cd keep && alembic revision --autogenerate -m "Your message"
+```
+
+Hint: make sure your models are imported at `./api/models/db/migrations/env.py` for autogenerator to pick them up.
+
 ## VSCode
 You can run Keep from your VSCode (after cloning the repo) by adding this configurations to your `.vscode/launch.json`:
 

keep/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 = 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 [%(name)s] %(message)s
+datefmt = %H:%M:%S

keep/api/config.py

@@ -4,7 +4,7 @@
 import keep.api.logging
 from keep.api.api import AUTH_TYPE
 from keep.api.core.config import AuthenticationType
-from keep.api.core.db_on_start import create_db_and_tables, try_create_single_tenant
+from keep.api.core.db_on_start import migrate_db, try_create_single_tenant
 from keep.api.core.dependencies import SINGLE_TENANT_UUID
 
 PORT = int(os.environ.get("PORT", 8080))
@@ -17,7 +17,7 @@ def on_starting(server=None):
     """This function is called by the gunicorn server when it starts"""
     logger.info("Keep server starting")
     if not os.environ.get("SKIP_DB_CREATION", "false") == "true":
-        create_db_and_tables()
+        migrate_db()
 
     # Create single tenant if it doesn't exist
     if AUTH_TYPE in [

keep/api/core/db_on_start.py

@@ -17,9 +17,11 @@
 import logging
 import os
 
+import alembic.command
+import alembic.config
+
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy_utils import create_database, database_exists
-from sqlmodel import Session, SQLModel, select
+from sqlmodel import Session, select
 
 from keep.api.core.db_utils import create_db_engine
 
@@ -50,7 +52,7 @@ def try_create_single_tenant(tenant_id: str) -> None:
             User,
         )
 
-        create_db_and_tables()
+        migrate_db()
     except Exception:
         pass
     with Session(engine) as session:
@@ -99,19 +101,16 @@ def try_create_single_tenant(tenant_id: str) -> None:
             pass
 
 
-def create_db_and_tables():
+
+def migrate_db():
     """
-    Creates the database and tables.
+    Run migrations to make sure the DB is up-to-date.
     """
-    try:
-        if not database_exists(engine.url):
-            logger.info("Creating the database")
-            create_database(engine.url)
-            logger.info("Database created")
-    # On Cloud Run, it fails to check if the database exists
-    except Exception:
-        logger.warning("Failed to create the database or detect if it exists.")
-        pass
-    logger.info("Creating the tables")
-    SQLModel.metadata.create_all(engine)
-    logger.info("Tables created")
+    logger.info("Running migrations...")
+    config_path = 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)
+    config.set_main_option("script_location", os.path.dirname(os.path.abspath(__file__)) + "/../models/db/migrations")
+    alembic.command.upgrade(config, "head")
+    logger.info("Finished migrations")

keep/api/models/db/migrations/env.py

@@ -0,0 +1,89 @@
+import asyncio
+from logging.config import fileConfig
+
+from alembic import context
+from sqlalchemy.future import Connection
+from sqlmodel import SQLModel
+
+from keep.api.core.db_utils import create_db_engine
+
+from keep.api.models.db.alert import *
+from keep.api.models.db.action import *
+from keep.api.models.db.dashboard import *
+from keep.api.models.db.extraction 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.tenant import *
+from keep.api.models.db.rule import *
+from keep.api.models.db.user import *
+from keep.api.models.db.workflow import *
+from keep.api.models.db.dashboard 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:
+    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_db_engine()
+    context.configure(
+        url=str(connectable.url),
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    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)
+
+    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()
+    do_run_migrations(connectable.connect())
+
+
+loop = asyncio.get_event_loop()
+if context.is_offline_mode():
+    task = run_migrations_offline()
+else:
+    task = run_migrations_online()
+
+loop.run_until_complete(task)

keep/api/models/db/migrations/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+
+${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"}