refactoring of collection

db/schema.sql

@@ -1,6 +1,6 @@
 CREATE TABLE vms (
     id INTEGER PRIMARY KEY,
-    uuid TEXT NOT NULL,
+    uuid TEXT NOT NULL UNIQUE,
     hostname TEXT NOT NULL,
     cores REAL NOT NULL,
     mem REAL NOT NULL,
@@ -16,9 +16,9 @@ CREATE TABLE builds (
     vm INTEGER NOT NULL,
     start INTEGER NOT NULL,
     end INTEGER NOT NULL,
-    job_id INTEGER NOT NULL,
+    job_id INTEGER NOT NULL UNIQUE,
     job_status TEXT NOT NULL,
-    num_retries INTEGER NOT NULL,
+    retries INTEGER NOT NULL,
     ref TEXT NOT NULL,
     pkg_name TEXT NOT NULL,
     pkg_version TEXT NOT NULL,

gantry/__main__.py

@@ -1,5 +1,24 @@
+import os
+
+import aiosqlite
+from aiohttp import web
+
+from gantry.views import routes
+
+
+async def init_db(app: web.Application):
+    db = await aiosqlite.connect(os.environ["DB_FILE"])
+    await db.execute("PRAGMA foreign_keys = ON;")
+    app["db"] = db
+    yield
+    await db.close()
+
+
 def main():
-    print("Hello World")
+    app = web.Application()
+    app.add_routes(routes)
+    app.cleanup_ctx.append(init_db)
+    web.run_app(app)
 
 
 if __name__ == "__main__":

gantry/collection.py

@@ -0,0 +1,91 @@
+import logging
+
+import aiosqlite
+
+from gantry.models import VM, Build
+from gantry.util.gitlab import GitlabClient
+from gantry.util.prometheus import IncompleteData, PrometheusClient
+
+
+async def fetch_build(payload: dict, db: aiosqlite.Connection) -> None:
+    """
+    Fetches a job's information from Prometheus and inserts it into the database.
+    If there is data missing at any point, the function will still return so the webhook
+    responds as expected. If an exception is thrown, that behavior was unanticipated by
+    this program and should be investigated.
+
+    args:
+        payload: a dictionary containing the information from the Gitlab job hook
+        db: an active aiosqlite connection
+
+    returns: None in order to accomodate a 200 response for the webhook.
+    """
+
+    gitlab = GitlabClient()
+    prometheus = PrometheusClient()
+
+    build = Build(
+        status=payload["build_status"],
+        name=payload["build_name"],
+        id=payload["build_id"],
+        start=payload["build_started_at"],
+        end=payload["build_finished_at"],
+        retries=payload["retries_count"],
+        ref=payload["ref"],
+    )
+
+    # perform checks to see if we should collect data for this job
+    if (
+        build.status not in ("success",)
+        or not build.valid_name  # is not a build job
+        or await build.in_db(db)  # job already in the database
+        or await build.is_ghost(db, gitlab)
+    ):
+        return
+
+    try:
+        await build.get_annotations(prometheus)
+        await build.get_resources(prometheus)
+        await build.get_usage(prometheus)
+        vm_id = await fetch_vm(db, prometheus, build.node, build.midpoint)
+    except IncompleteData as e:
+        # missing data, skip this job
+        logging.error(e)
+        return
+
+    await build.insert(db, vm_id)
+    # vm and build will get saved at the same time to make sure
+    # we don't accidentally commit a vm without a build
+    await db.commit()
+
+    return
+
+
+async def fetch_vm(
+    db: aiosqlite.Connection,
+    prometheus: PrometheusClient,
+    hostname: dict,
+    query_time: float,
+) -> int:
+    """
+    Finds an existing VM in the database or inserts a new one.
+
+    args:
+        db: an active aiosqlite connection
+        prometheus:
+        hostname: the hostname of the VM
+        query_time: any point during VM runtime, usually grabbed from build
+
+    returns: id of the inserted or existing VM
+    """
+    vm = VM(
+        hostname=hostname,
+        query_time=query_time,
+    )
+
+    # do not proceed if the VM exists
+    if existing_vm := await vm.db_id(db, prometheus):
+        return existing_vm
+
+    await vm.get_labels(prometheus)
+    return await vm.insert(db)

gantry/models/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa
+from .build import Build
+from .vm import VM