Merge pull request #319 from doaa-altarawy/sql_migration

Add upgrade DB to CLI

MANIFEST.in

@@ -12,5 +12,6 @@ include MANIFEST.in
 
 include versioneer.py
 include qcfractal/_version.py
+include qcfractal/alembic.ini
 
 prune benchmarks

devtools/scripts/create_staging.py

@@ -120,14 +120,17 @@ def copy_users(staging_storage, prod_storage):
     print('---- Done copying Users\n\n')
 
 
-def copy_managers(staging_storage, prod_storage):
+def copy_managers(staging_storage, prod_storage, mang_list):
     """Copy ALL managers from prod to staging"""
 
-    prod_mangers = prod_storage.get_managers()
-    print('-----Total # of Managers in the DB is: ', prod_mangers['meta']['n_found'])
+    prod_mangers = []
+    for mang in mang_list:
+        prod_mangers.extend(prod_storage.get_managers(name=mang)['data'])
 
+    print('-----Total # of Managers to copy is: ', len(prod_mangers))
 
-    sql_insered = staging_storage._copy_managers(prod_mangers['data'])['data']
+
+    sql_insered = staging_storage._copy_managers(prod_mangers)['data']
     print('Inserted in SQL:', len(sql_insered))
 
     print('---- Done copying Queue Manager\n\n')
@@ -408,7 +411,7 @@ def copy_task_queue(staging_storage, production_storage, SAMPLE_SIZE=None):
     print('Copying {} TaskQueues'.format(count_to_copy))
 
 
-    base_results = []
+    base_results, managers = [], []
     results = {
         'result': [],
         'optimization_procedure': [],
@@ -418,6 +421,7 @@ def copy_task_queue(staging_storage, production_storage, SAMPLE_SIZE=None):
 
     for rec in prod_tasks:
         base_results.append(rec.base_result.id)
+        managers.append(rec.manager)
 
     with production_storage.session_scope() as session:
         ret = session.query(BaseResultORM.id, BaseResultORM.result_type).filter(BaseResultORM.id.in_(base_results)).all()
@@ -431,7 +435,7 @@ def copy_task_queue(staging_storage, production_storage, SAMPLE_SIZE=None):
     proc_map1 = copy_optimization_procedure(staging_storage, production_storage, procedure_ids=results['optimization_procedure'])
     proc_map2 = copy_grid_optimization_procedure(staging_storage, production_storage, procedure_ids=results['grid_optimization_procedure'])
     proc_map3 = copy_torsiondrive_procedure(staging_storage, production_storage, procedure_ids=results['torsiondrive_procedure'])
-
+    copy_managers(staging_storage, production_storage, managers)
 
     for rec in prod_tasks:
         id = int(rec.base_result.id)
@@ -465,9 +469,8 @@ def main():
         print('Exit without creating the DB.')
         return
 
-    # copy all managers and users, small tables, no need for sampling
+    # copy all users, small tables, no need for sampling
     copy_users(staging_storage, production_storage)
-    copy_managers(staging_storage, production_storage)
 
     # copy sample of results and procedures
     print('\n-------------- Results -----------------')

qcfractal/alembic.ini

@@ -2,7 +2,7 @@
 
 [alembic]
 # path to migration scripts
-script_location = alembic
+script_location = qcfractal:alembic
 
 # template used to generate migration files
 # file_template = %%(rev)s_%%(slug)s
@@ -35,7 +35,8 @@ script_location = alembic
 # are written from script.py.mako
 # output_encoding = utf-8
 
-sqlalchemy.url = postgresql+psycopg2://qcarchive:mypass@localhost:5432/staging_qcarchivedb
+# NOT used from here, will be read from FractalConfig
+sqlalchemy.url = 'Overridden by the application'
 
 
 # Logging configuration

qcfractal/alembic/env.py

@@ -8,6 +8,9 @@
 
 from qcfractal.storage_sockets import sql_models
 
+import yaml
+from qcfractal.config import FractalConfig
+
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config = context.config
@@ -29,6 +32,10 @@
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
+# Overwrite the ini-file sqlalchemy.url path
+uri = context.get_x_argument(as_dictionary=True).get('uri')
+config.set_main_option('sqlalchemy.url', uri)
+
 
 def run_migrations_offline():
     """Run migrations in 'offline' mode.

qcfractal/cli/qcfractal_server.py

@@ -31,7 +31,8 @@ def parse_args():
         cli_name = "--" + field.replace("_", "-")
         server_init.add_argument(cli_name, **FractalServerSettings.help_info(field))
 
-    init.add_argument("--overwrite", action='store_true', help="Overwrites the current configuration file.")
+    init.add_argument("--overwrite-config", action='store_true', help="Overwrites the current configuration file.")
+    init.add_argument("--clear-database", default='store_true', help="Clear the content of the given database and initialize it.")
     init.add_argument("--base-folder", **FractalConfig.help_info("base_folder"))
 
     ### Start subcommands
@@ -44,10 +45,6 @@ def parse_args():
         cli_name = "--" + field.replace("_", "-")
         fractal_args.add_argument(cli_name, **FractalServerSettings.help_info(field))
 
-    fractal_args.add_argument("--database-name",
-                              default=None,
-                              type=str,
-                              help="The database to connect to, defaults to the default database name.")
     fractal_args.add_argument("--server-name", **FractalServerSettings.help_info("name"))
     fractal_args.add_argument(
         "--start-periodics",
@@ -63,6 +60,10 @@ def parse_args():
     fractal_args.add_argument("--tls-cert", type=str, default=None, help="Certificate file for TLS (in PEM format)")
     fractal_args.add_argument("--tls-key", type=str, default=None, help="Private key file for TLS (in PEM format)")
 
+    ### Upgrade subcommands
+    upgrade = subparsers.add_parser('upgrade', help="Upgrade QCFractal database.")
+    upgrade.add_argument("--base-folder", **FractalConfig.help_info("base_folder"))
+
     compute_args = start.add_argument_group('Local Computation Settings')
     compute_args.add_argument("--local-manager",
                               const=-1,
@@ -112,20 +113,22 @@ def parse_args():
 
 
 def server_init(args, config):
+    # alembic stamp head
 
     print("Initializing QCFractal configuration.")
     # Configuration settings
 
     config.base_path.mkdir(parents=True, exist_ok=True)
-    overwrite = args.get("overwrite", False)
+    overwrite_config = args.get("overwrite_config", False)
+    clear_database = args.get("clear_database", False)
 
     psql = PostgresHarness(config, quiet=False, logger=print)
 
     # Make sure we do not delete anything.
     if config.config_file_path.exists():
         print()
-        if not overwrite:
-            print("QCFractal configuration file already exists, to overwrite use '--overwrite' "
+        if not overwrite_config:
+            print("QCFractal configuration file already exists, to overwrite use '--overwrite-config' "
                   "or use the `qcfractal-server config` command line to alter settings.")
             sys.exit(2)
         else:
@@ -159,11 +162,14 @@ def server_init(args, config):
     print("\n>>> Settings found:\n")
     print(print_config)
 
+    print("\n>>> Writing settings...")
+    config.config_file_path.write_text(yaml.dump(config.dict(), default_flow_style=False))
+
     print("\n>>> Setting up PostgreSQL...\n")
     config.database_path.mkdir(exist_ok=True)
     if config.database.own:
         try:
-            psql.initialize()
+            psql.initialize_postgres()
         except ValueError as e:
             print(str(e))
             sys.exit(1)
@@ -172,9 +178,15 @@ def server_init(args, config):
             "Own was set to False, QCFractal will expect a live PostgreSQL server with the above connection information."
         )
 
-    print("\n>>> Writing settings...")
-    config.config_file_path.write_text(yaml.dump(config.dict(), default_flow_style=False))
+    if config.database.own or clear_database:
+        print("\n>>> Initializing database schema...\n")
+        try:
+            psql.init_database()
+        except ValueError as e:
+            print(str(e))
+            sys.exit(1)
 
+    # create tables and stamp version (if not)
     print("\n>>> Finishing up...")
     print("\n>>> Success! Please run `qcfractal-server start` to boot a FractalServer!")
 
@@ -186,6 +198,7 @@ def server_config(args, config):
 
 
 def server_start(args, config):
+    # check if db not current, ask for upgrade
 
     print("Starting a QCFractal server.\n")
 
@@ -226,7 +239,6 @@ def server_start(args, config):
         logfile = str(config.base_path / config.fractal.logfile)
 
     print("\n>>> Checking the PostgreSQL connection...")
-    database_name = args.get("database_name", None) or config.database.default_database
     psql = PostgresHarness(config, quiet=False, logger=print)
 
     if not psql.is_alive():
@@ -237,7 +249,8 @@ def server_start(args, config):
             print(str(e))
             sys.exit(1)
 
-    psql.create_database(database_name)
+    # make sure DB is created
+    psql.create_database(config.database.database_name)
 
     print("\n>>> Initializing the QCFractal server...")
     try:
@@ -253,7 +266,7 @@ def server_start(args, config):
 
             # Database
             storage_uri=config.database_uri(safe=False, database=""),
-            storage_project_name=database_name,
+            storage_project_name=config.database.database_name,
             query_limit=config.fractal.query_limit,
 
             # Log options
@@ -280,6 +293,32 @@ def server_start(args, config):
     print("\n>>> Starting the QCFractal server...")
     server.start(start_periodics=args["start_periodics"])
 
+def server_upgrade(args, config):
+    # alembic upgrade head
+
+    print("Upgrading QCFractal server.\n")
+
+    print(f"QCFractal server base folder: {config.base_folder}")
+
+    print("\n>>> Checking the PostgreSQL connection...")
+    psql = PostgresHarness(config, quiet=False, logger=print)
+
+    if not psql.is_alive():
+        try:
+            print("\nCould not detect a PostgreSQL from configuration options, starting a PostgreSQL server.\n")
+            psql.start()
+        except ValueError as e:
+            print(str(e))
+            sys.exit(1)
+
+    print("\n>>> Upgrading the Database...")
+
+    try:
+        psql.upgrade()
+    except ValueError as e:
+        print(str(e))
+        sys.exit(1)
+
 
 def main(args=None):
 
@@ -307,7 +346,8 @@ def main(args=None):
             print(f"Could not find configuration file: {config.config_file_path}")
             sys.exit(1)
 
-        file_dict = FractalConfig(**yaml.load(config.config_file_path.read_text())).dict()
+        file_dict = FractalConfig(**yaml.load(config.config_file_path.read_text(),
+                                  Loader=yaml.FullLoader)).dict()
         config_dict = config.dict(skip_defaults=True)
 
         # Only fractal options can be changed by user input parameters
@@ -321,6 +361,8 @@ def main(args=None):
         server_config(args, config)
     elif command == "start":
         server_start(args, config)
+    elif command == 'upgrade':
+        server_upgrade(args, config)
 
 
 if __name__ == '__main__':

qcfractal/cli/tests/test_cli.py

@@ -23,7 +23,8 @@ def qcfractal_base_init(postgres_server):
 
     args = [
         "qcfractal-server", "init", "--base-folder",
-        str(tmpdir.name), "--db-own=False", f"--db-port={postgres_server.config.database.port}"
+        str(tmpdir.name), "--db-own=False", "--clear-database=True",
+        f"--db-port={postgres_server.config.database.port}"
     ]
     assert testing.run_process(args, **_options)
 
@@ -33,7 +34,12 @@ def qcfractal_base_init(postgres_server):
 @testing.mark_slow
 def test_cli_server_boot(qcfractal_base_init):
     port = "--port=" + str(testing.find_open_port())
-    args = ["qcfractal-server", "start", qcfractal_base_init, "--database-name=test_cli_db", port]
+    args = ["qcfractal-server", "start", qcfractal_base_init, port]
+    assert testing.run_process(args, interupt_after=10, **_options)
+
+@testing.mark_slow
+def test_cli_upgrade(qcfractal_base_init):
+    args = ["qcfractal-server", "upgrade", qcfractal_base_init]
     assert testing.run_process(args, interupt_after=10, **_options)
 
 

qcfractal/config.py

@@ -8,6 +8,7 @@
 from typing import Optional
 
 from pydantic import BaseSettings, Schema, validator
+import yaml
 
 from .util import doc_formatter
 
@@ -64,7 +65,7 @@ class DatabaseSettings(ConfigSettings):
     password: str = Schema(None, description="The postgres password for the give user.")
     directory: str = Schema(
         None, description="The physical location of the QCFractal instance data, defaults to the root folder.")
-    default_database: str = Schema("qcfractal_default", description="The default database to connect to.")
+    database_name: str = Schema("qcfractal_default", description="The database name to connect to.")
     logfile: str = Schema("qcfractal_postgres.log", description="The logfile to write postgres logs.")
     own: bool = Schema(True, description="If own is True, QCFractal will control the database instance. If False "
                                          "Postgres will expect a booted server at the database specification.")
@@ -114,6 +115,9 @@ class FractalConfig(ConfigSettings):
     Top level configuration headers and options for a QCFractal Configuration File
     """
 
+    # class variable, not in the pydantic model
+    _defaults_file_path: str = os.path.expanduser("~/.qca/qcfractal_defaults.yaml")
+
     base_folder: str = Schema(os.path.expanduser("~/.qca/qcfractal"),
                               description="The QCFractal base instance to attach to. "
                                           "Default will be your home directory")
@@ -123,6 +127,24 @@ class FractalConfig(ConfigSettings):
     class Config(SettingsCommonConfig):
         pass
 
+    def __init__(self, **kwargs):
+
+        # If no base_folder provided, read it from ~/.qca/qcfractal_defaults.yaml (if it exists)
+        # else, use the default base_folder
+        if 'base_folder' not in kwargs:
+            if Path(FractalConfig._defaults_file_path).exists():
+                with open(FractalConfig._defaults_file_path, "r") as handle:
+                    kwargs['base_folder'] = yaml.load(handle.read(),
+                                                      Loader=yaml.FullLoader)['default_base_folder']
+
+        super().__init__(**kwargs)
+
+    @classmethod
+    def from_base_folder(cls, base_folder):
+        path = Path(base_folder) / "qcfractal_config.yaml"
+        with open(str(path), "r") as handle:
+            return cls(**yaml.load(handle.read(), Loader=yaml.FullLoader))
+
     @property
     def base_path(self):
         return Path(self.base_folder)
@@ -156,7 +178,7 @@ def database_uri(self, safe=True, database=None):
         uri += f"{self.database.host}:{self.database.port}/"
 
         if database is None:
-            uri += self.database.default_database
+            uri += self.database.database_name
         else:
             uri += database