Simplify the code, remove unnecessary fragments

MANIFEST.in

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

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

@@ -24,13 +24,7 @@
 
 from .config import app_config
 from .util import mquery_version
-from .db import (
-    Database,
-    is_alembic_version_present,
-    is_database_initilized,
-    run_alembic_legacy_stamp,
-    run_alembic_upgrade,
-)
+from .db import Database
 from .lib.yaraparse import parse_yara
 from .plugins import PluginManager
 from .lib.ursadb import UrsaDb
@@ -54,21 +48,7 @@
 
 @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()
+    db.alembic_upgrade()
     yield
 
 

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,13 +12,11 @@
 from rq import Queue  # type: ignore
 from sqlmodel import (
     Session,
-    SQLModel,
     create_engine,
     select,
     and_,
     update,
     col,
-    inspect,
 )
 
 from .models.agentgroup import AgentGroup
@@ -30,9 +31,6 @@
 # "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"
@@ -414,51 +412,7 @@ def set_config_key(self, plugin_name: str, key: str, value: str) -> None:
             session.add(entry)
             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)
-
-
-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")