Initialize the database automatically when starting mquery #402 (#421)

Automatically init database with alembic upgrade head

Co-authored-by: Michał Jura <[email protected]>
Co-authored-by: msm <[email protected]>

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.

MANIFEST.in

@@ -1 +1,2 @@
 graft src/mqueryfront/dist
+include src/alembic.ini

README.md

@@ -25,7 +25,6 @@ git clone https://github.com/CERT-Polska/mquery.git
 cd mquery
 vim .env  # optional - change samples and index directory locations
 docker compose up --scale daemon=3  # building the images will take a while
-docker compose exec web python3 -m mquery.db
 ```
 
 The web interface should be available at `http://localhost`.

docs/how-to/install-native.md

@@ -118,7 +118,6 @@ Now you need to create and configure a database
 ```shell
 psql -c "CREATE DATABASE mquery"
 source /opt/mquery/venv/bin/activate  # remember, we need virtualenv
-python3 -m mquery.db  # initialize the mquery database
 ```
 
 ### Start everything

setup.py

@@ -11,6 +11,8 @@
         "mquery.lib",
         "mquery.plugins",
         "mquery.models",
+        "mquery.migrations",
+        "mquery.migrations.versions",
     ],
     package_dir={"mquery": "src"},
     include_package_data=True,

src/alembic.ini

@@ -1,5 +1,5 @@
 [alembic]
-script_location = migrations
+script_location = %(here)s/migrations
 prepend_sys_path = .
 version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
 

src/app.py

@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 import os
 
 import uvicorn  # type: ignore
@@ -44,8 +45,15 @@
     ServerSchema,
 )
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    db.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

@@ -1,3 +1,6 @@
+from alembic.config import Config
+from alembic import command
+from pathlib import Path
 from collections import defaultdict
 from contextlib import contextmanager
 from typing import List, Optional, Dict, Any
@@ -9,7 +12,6 @@
 from rq import Queue  # type: ignore
 from sqlmodel import (
     Session,
-    SQLModel,
     create_engine,
     select,
     and_,
@@ -410,11 +412,7 @@ def set_config_key(self, plugin_name: str, key: str, value: str) -> None:
             session.add(entry)
             session.commit()
 
-
-def init_db() -> None:
-    engine = create_engine(app_config.database.url, echo=True)
-    SQLModel.metadata.create_all(engine)
-
-
-if __name__ == "__main__":
-    init_db()
+    def alembic_upgrade(self) -> None:
+        config_file = Path(__file__).parent / "alembic.ini"
+        alembic_cfg = Config(str(config_file))
+        command.upgrade(alembic_cfg, "head")