Build state tracker

Define a SQLAlchemy DB class "Dataset" which defines persistent state for a
Pbench dataset including most importantly the dataset's username (which is not
otherwise recorded until we index) and the dataset's current state so that we
can track the progress of a dataset through the Pbench server pipeline.

We also support Metadata associated with each Dataset, describing additional
information about datasets beyond the "state". For example, the backup
component will mark a dataset as "ARCHIVED" and pbench-reindex marks the
selected datasets to "REINDEX".

A new dataset starts in UPLOADING state and will progress through the steps as
we perform operations on it. We're defining both "in progress" -ing steps such
as INDEXING as well as completion steps (which are ready for the next
operation) such as INDEXED. There are also two "terminal" states, EXPIRED and
QUARANTINED, from which a dataset cannot exit.

NOTE that the SqlAlchemy infrastructure borrows heavily from Nikhil's work on
the user model; however I tried to improve initialization. I'm specifying our
own database, which we expect to be part of the Postgres URI in the config
file. My assumption and intent is that this name will resemble the ES index
prefix (although Postgres doesn't like '-' characters); so, for example, with
an index prefix of "pbench-drb" I'd choose a DB name like "pbench_drb". This
allows multiple Pbench server instances to coexist on a single Postgres
server. I also added some logic to create the DB if necessary and to create
the DB tables before starting gunicorn ... as Postgres can have problems when
multiple clients try to do this at the same time.

Author: dbutenhof

.gitignore

@@ -7,6 +7,7 @@ build
 dist
 *~
 ????-*.patch
+.env
 .npmrc
 .yarnrc
 *.egg-info

jenkins/development.Dockerfile

@@ -97,6 +97,7 @@ RUN \
         rsync \
         screen \
         sos \
+        sqlite \
         tar \
         xz \
         && \

lib/pbench/cli/server/shell.py

@@ -4,9 +4,12 @@
 import sys
 
 from configparser import NoSectionError, NoOptionError
+from sqlalchemy_utils import database_exists, create_database
 
 from pbench.common.exceptions import BadConfig, ConfigFileNotSpecified
 from pbench.server.api import create_app, get_server_config
+from pbench.server.db.database import Database
+from pbench.common.logger import get_pbench_logger
 
 
 def app():
@@ -24,12 +27,24 @@ def main():
     except (ConfigFileNotSpecified, BadConfig) as e:
         print(e)
         sys.exit(1)
+    logger = get_pbench_logger(__name__, server_config)
     try:
         host = str(server_config.get("pbench-server", "bind_host"))
         port = str(server_config.get("pbench-server", "bind_port"))
+        db = str(server_config.get("Postgres", "db_uri"))
         workers = str(server_config.get("pbench-server", "workers"))
-    except (NoOptionError, NoSectionError) as e:
-        print(f"{__name__}: ERROR: {e.__traceback__}")
+
+        # Multiple gunicorn workers will attempt to connect to the DB; rather
+        # than attempt to synchronize them, detect a missing DB (from the
+        # postgres URI) and create it here. It's safer to do this here,
+        # where we're single-threaded.
+        if not database_exists(db):
+            logger.info("Postgres DB {} doesn't exist", db)
+            create_database(db)
+            logger.info("Created DB {}", db)
+        Database.init_db(server_config, logger)
+    except (NoOptionError, NoSectionError):
+        logger.exception(f"{__name__}: ERROR")
         sys.exit(1)
 
     subprocess.run(

lib/pbench/server/api/__init__.py

@@ -21,7 +21,6 @@
 from pbench.server.database.database import Database
 from pbench.server.api.resources.query_apis.query_month_indices import QueryMonthIndices
 from pbench.server.api.auth import Auth
-
 from pbench.server.api.resources.users_api import (
     RegisterUser,
     Login,

lib/pbench/server/api/resources/__init__.py

@@ -7,6 +7,8 @@
 from flask import request, jsonify
 from werkzeug.utils import secure_filename
 from pbench.server.utils import filesize_bytes
+from pbench.server.db.models.tracker import Dataset, States
+
 
 ALLOWED_EXTENSIONS = {"xz"}
 
@@ -109,6 +111,12 @@ def put(self, controller):
         md5_full_path = Path(path, f"{filename}.md5")
         bytes_received = 0
 
+        # TODO: Need real user from PUT!
+
+        # Create a tracking dataset object; it'll begin in UPLOADING state
+        dataset = Dataset.create(controller=controller, path=tar_full_path, md5=md5sum)
+        self.logger.info("Uploading {}", dataset)
+
         with tempfile.NamedTemporaryFile(mode="wb", dir=path) as ofp:
             chunk_size = 4096
             self.logger.debug("Writing chunks")
@@ -181,6 +189,11 @@ def put(self, controller):
                 )
                 raise
 
+        if dataset:
+            try:
+                dataset.advance(States.UPLOADED)
+            except Exception:
+                self.logger.exception("Unable to finalize {}", dataset)
         response = jsonify(dict(message="File successfully uploaded"))
         response.status_code = 201
         return response

lib/pbench/server/api/resources/upload_api.py

@@ -0,0 +1,61 @@
+import sys
+from sqlalchemy import create_engine
+from sqlalchemy.orm import scoped_session, sessionmaker
+from sqlalchemy.ext.declarative import declarative_base
+from pbench.server import NoOptionError, NoSectionError
+
+
+class Database:
+    """ Database - a base class for declaring SQLAlchemy database model
+    classes.
+    """
+
+    Base = declarative_base()
+
+    # Initialize the db scoped session
+    db_session = scoped_session(sessionmaker(autocommit=False, autoflush=False))
+
+    @staticmethod
+    def get_engine_uri(config, logger):
+        try:
+            psql = config.get("Postgres", "db_uri")
+            return psql
+        except (NoSectionError, NoOptionError):
+            if logger:
+                logger.error("Failed to find [Postgres] db_uri in configuration file.")
+            sys.exit(1)
+
+        # return f"postgresql+{psql_driver}://{psql_username}:{psql_password}@{psql_host}:{psql_port}/{psql_db}"
+
+    @staticmethod
+    def init_engine(server_config, logger):
+        try:
+            uri = Database.get_engine_uri(server_config, logger)
+            return create_engine(uri)
+        except Exception:
+            if logger:
+                logger.exception(
+                    "Exception while creating a sqlalchemy engine for {}", uri
+                )
+            sys.exit(1)
+
+    @staticmethod
+    def init_db(server_config, logger):
+        # Attach the logger to the base class for models to find
+        if not hasattr(Database.Base, "logger"):
+            Database.Base.logger = logger
+
+        # IMPORTANT:
+        # Make sure all the models are imported before this function is called
+        # so that they will be registered properly on the metadata. Otherwise
+        # metadata will not have any tables and create_all functionality will
+        # do nothing.
+        try:
+            Database.Base.query = Database.db_session.query_property()
+            engine = Database.init_engine(server_config, logger)
+            Database.Base.metadata.create_all(bind=engine)
+            Database.db_session.configure(bind=engine)
+        except Exception:
+            if logger:
+                logger.exception("Exception while initializing sqlalchemy database")
+            sys.exit(1)