Initialize the database automatically when starting mquery

INSTALL.md

@@ -20,7 +20,6 @@ mkdir samples
 # expect files in ./samples directory, and keep index in ./index.
 vim .env
 docker compose up --scale daemon=3  # this will take a while
-docker compose exec web python3 -m mquery.db
 ```
 
 - Good for testing mquery and production deployments on a single server
@@ -38,7 +37,6 @@ cd mquery
 # expect files in ./samples directory, and keep index in ./index.
 vim .env
 docker compose -f docker-compose.dev.yml up  # this will take a while
-docker compose exec dev-web python3 -m mquery.db
 ```
 
 - Good for development - all file changes will be picked up automatically.

src/app.py

@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 import os
 
 import uvicorn  # type: ignore
@@ -23,7 +24,13 @@
 
 from .config import app_config
 from .util import mquery_version
-from .db import Database
+from .db import (
+    Database,
+    is_alembic_version_present,
+    is_database_initilized,
+    run_alembic_legacy_stamp,
+    run_alembic_upgrade,
+)
 from .lib.yaraparse import parse_yara
 from .plugins import PluginManager
 from .lib.ursadb import UrsaDb
@@ -44,8 +51,29 @@
     ServerSchema,
 )
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Check if the database is initilized
+    if is_database_initilized():
+        # Check if database contains alembic_version table
+        if is_alembic_version_present():
+            # If True it means that we can run alembic upgrade head without worry.
+            pass
+        # If False database is not up-to-date
+        # but we can't just run alembic head upgrade as there is no alembic_version table
+        # so we need to run alemibc stamp dbb81bd4d47f
+        # as dbb81bd4d47f is the number of last migration before alembic head upgrade was added
+        else:
+            run_alembic_legacy_stamp()
+
+    # finally we can run alembic upgrade head to upgrade (if needed) the database
+    run_alembic_upgrade()
+    yield
+
+
 db = Database(app_config.redis.host, app_config.redis.port)
-app = FastAPI()
+app = FastAPI(lifespan=lifespan)
 
 
 def with_plugins() -> Iterable[PluginManager]:

src/db.py

@@ -15,6 +15,7 @@
     and_,
     update,
     col,
+    inspect,
 )
 
 from .models.agentgroup import AgentGroup
@@ -29,6 +30,9 @@
 # "Magic" plugin name, used for configuration of mquery itself
 MQUERY_PLUGIN_NAME = "Mquery"
 
+# Legacy purpose
+LAST_REVISION_BEFORE_ALEMBiC_HEAD_UPGRADE = "dbb81bd4d47f"
+
 
 class TaskType(Enum):
     SEARCH = "search"
@@ -411,6 +415,46 @@ def set_config_key(self, plugin_name: str, key: str, value: str) -> None:
             session.commit()
 
 
+def is_alembic_version_present() -> bool:
+    engine = create_engine(app_config.database.url, echo=True)
+    with engine.connect() as connection:
+        inspector = inspect(connection)
+        table_names = inspector.get_table_names()
+        if "alembic_version" in table_names:
+            return True
+    return False
+
+
+def is_database_initilized() -> bool:
+    engine = create_engine(app_config.database.url, echo=True)
+    with engine.connect() as connection:
+        inspector = inspect(connection)
+        table_names = inspector.get_table_names()
+        if table_names:
+            return True
+    return False
+
+
+def run_alembic_legacy_stamp():
+    from alembic.config import Config
+    from alembic import command
+    from os import chdir
+
+    chdir("/usr/src/app/src/")
+    alembic_cfg = Config("alembic.ini")
+    command.stamp(alembic_cfg, LAST_REVISION_BEFORE_ALEMBiC_HEAD_UPGRADE)
+
+
+def run_alembic_upgrade() -> None:
+    from alembic.config import Config
+    from alembic import command
+    from os import chdir
+
+    chdir("/usr/src/app/src/")
+    alembic_cfg = Config("alembic.ini")
+    command.upgrade(alembic_cfg, "head")
+
+
 def init_db() -> None:
     engine = create_engine(app_config.database.url, echo=True)
     SQLModel.metadata.create_all(engine)