Merge branch 'develop' into Issue955/fix-cancel-workflows

beeflow/client/core.py

@@ -18,6 +18,7 @@
 import datetime
 import time
 import importlib.metadata
+from pathlib import Path
 import packaging
 
 import daemon
@@ -26,6 +27,7 @@
 
 from beeflow.client import bee_client
 from beeflow.common.config_driver import BeeConfig as bc
+from beeflow.common.config_driver import AlterConfig
 from beeflow.common import cli_connection
 from beeflow.common import paths
 from beeflow.wf_manager.resources import wf_utils
@@ -35,6 +37,9 @@
 from beeflow.common.deps import redis_manager
 
 
+REPO_PATH = Path(*Path(__file__).parts[:-3])
+
+
 class ComponentManager:
     """Component manager class."""
 
@@ -621,19 +626,26 @@ def pull_to_tar(ref, tarball):
     subprocess.check_call(['ch-convert', '-i', 'ch-image', '-o', 'tar', ref, tarball])
 
 
+def build_to_tar(tag, dockerfile, tarball):
+    """Build a container from a Dockerfile and convert to tarball."""
+    subprocess.check_call(['ch-image', 'build', '-t', tag, '-f', dockerfile, '.'])
+    subprocess.check_call(['ch-convert', '-i', 'ch-image', '-o', 'tar', tag, tarball])
+
+
 @app.command()
 def pull_deps(outdir: str = typer.Option('.', '--outdir', '-o',
                                          help='directory to store containers in')):
     """Pull required BEE containers and store in outdir."""
     load_check_charliecloud()
     neo4j_path = os.path.join(os.path.realpath(outdir), 'neo4j.tar.gz')
-    pull_to_tar('neo4j:5.17', neo4j_path)
+    neo4j_dockerfile = str(Path(REPO_PATH, "beeflow/data/dockerfiles/Dockerfile.neo4j"))
+    build_to_tar('neo4j_image', neo4j_dockerfile, neo4j_path)
     redis_path = os.path.join(os.path.realpath(outdir), 'redis.tar.gz')
     pull_to_tar('redis', redis_path)
-    print()
-    print('The BEE dependency containers have been successfully downloaded. '
-          'Please make sure to set the following options in your config:')
-    print()
-    print('[DEFAULT]')
-    print('neo4j_image =', neo4j_path)
-    print('redis_image =', redis_path)
+
+    AlterConfig(changes={'DEFAULT': {'neo4j_image': neo4j_path,
+                                     'redis_image': redis_path}}).save()
+
+    dep_dir = container_manager.get_dep_dir()
+    if os.path.isdir(dep_dir):
+        shutil.rmtree(dep_dir)

beeflow/common/config_driver.py

@@ -3,14 +3,14 @@
 from configparser import ConfigParser
 import getpass
 import os
-import base64
 import platform
 import random
 import shutil
 import textwrap
 import typer
 
 from beeflow.common.config_validator import ConfigValidator
+from beeflow.common import config_utils
 from beeflow.common.cli import NaturalOrderGroup
 from beeflow.common import validation
 from beeflow.common.tab_completion import filepath_completion
@@ -107,12 +107,7 @@ def init(cls, userconfig=None, **_kwargs):
                 print("Configuration file is missing! Generating new config file.")
                 new(USERCONFIG_FILE)
         # remove default keys from the other sections
-        default_keys = list(config['DEFAULT'])
-        config = {sec_name: {key: config[sec_name][key] for key in config[sec_name]
-                             if sec_name == 'DEFAULT' or key not in default_keys} # noqa
-                  for sec_name in config}
-        # Validate the config
-        cls.CONFIG = VALIDATOR.validate(config)
+        cls.CONFIG = config_utils.filter_and_validate(config, VALIDATOR)
 
     @classmethod
     def userconfig_path(cls):
@@ -323,12 +318,7 @@ def validate_chrun_opts(opts):
 VALIDATOR.option('graphdb', 'hostname', default='localhost',
                  info='hostname of database')
 
-
-# Generate random initial password for neo4j
-random_bytes = os.urandom(32)
-random_pass = base64.b64encode(random_bytes).decode('utf-8')
-
-VALIDATOR.option('graphdb', 'dbpass', default=random_pass, info='password for database')
+VALIDATOR.option('graphdb', 'dbpass', default='password', info='password for database')
 
 VALIDATOR.option('graphdb', 'gdb_image_mntdir', default=join_path('/tmp', USER),
                  info='graph database image mount directory', validator=validation.make_dir)
@@ -457,18 +447,7 @@ def save(self):
         if ans.lower() != 'y':
             print('Quitting without saving')
             return
-        try:
-            with open(self.fname, 'w', encoding='utf-8') as fp:
-                print('# BEE Configuration File', file=fp)
-                for sec_name, section in self.sections.items():
-                    if not section:
-                        continue
-                    print(file=fp)
-                    print(f'[{sec_name}]'.format(sec_name), file=fp)
-                    for opt_name in section:
-                        print(f'{opt_name} = {section[opt_name]}', file=fp)
-        except FileNotFoundError:
-            print('Configuration file does not exist!')
+        config_utils.write_config(self.fname, self.sections)
         print(70 * '#')
         print('Before running BEE, check defaults in the configuration file:',
               f'\n\t{self.fname}',
@@ -478,6 +457,78 @@ def save(self):
         print(70 * '#')
 
 
+class AlterConfig:
+    r"""Class to alter an existing BEE configuration.
+
+    Changes can be made when the class is instantiated, for example:
+    AlterConfig(changes={'DEFAULT': {'neo4j_image': '/path/to/neo4j'}})
+
+    Changes can also be made later, for example:
+    alter_config = AlterConfig()
+    alter_config.change_value('DEFAULT', 'neo4j_image', '/path/to/neo4j')
+    """
+
+    def __init__(self, fname=USERCONFIG_FILE, validator=VALIDATOR, changes=None):
+        """Load the existing configuration."""
+        self.fname = fname
+        self.validator = validator
+        self.config = None
+        self.changes = changes if changes is not None else {}
+        self._load_config()
+
+        for sec_name, opts in self.changes.items():
+            for opt_name, new_value in opts.items():
+                self.change_value(sec_name, opt_name, new_value)
+
+    def _load_config(self):
+        """Load the existing configuration file into memory."""
+        config = ConfigParser()
+        try:
+            with open(self.fname, encoding='utf-8') as fp:
+                config.read_file(fp)
+                self.config = config_utils.filter_and_validate(config, self.validator)
+        except FileNotFoundError:
+            for section_change in self.changes:
+                for option_change in self.changes[section_change]:
+                    for opt_name, option in VALIDATOR.options(section_change):
+                        if opt_name == option_change:
+                            option.default = self.changes[section_change][option_change]
+            self.config = ConfigGenerator(self.fname, self.validator).choose_values().sections
+
+    def change_value(self, sec_name, opt_name, new_value):
+        """Change the value of a configuration option."""
+        if sec_name not in self.config:
+            raise ValueError(f'Section {sec_name} not found in the config.')
+        if opt_name not in self.config[sec_name]:
+            raise ValueError(f'Option {opt_name} not found in section {sec_name}.')
+
+        # Find the correct option from the validator
+        options = self.validator.options(sec_name)  # Get all options for the section
+        for option_name, option in options:
+            if option_name == opt_name:
+                # Validate the new value before changing
+                option.validate(new_value)
+                self.config[sec_name][opt_name] = new_value
+                # Track changes in attribute
+                if sec_name not in self.changes:
+                    self.changes[sec_name] = {}
+                self.changes[sec_name][opt_name] = new_value
+                return
+
+        raise ValueError(f'Option {opt_name} not found in the validator for section {sec_name}.')
+
+    def save(self):
+        """Save the modified configuration back to the file."""
+        if os.path.exists(self.fname):
+            config_utils.backup(self.fname)
+        config_utils.write_config(self.fname, self.config)
+        # Print out changes
+        print("Configuration saved. The following values were changed:")
+        for sec_name, options in self.changes.items():
+            for opt_name, new_value in options.items():
+                print(f'Section [{sec_name}], Option [{opt_name}] changed to [{new_value}].')
+
+
 app = typer.Typer(no_args_is_help=True, add_completion=False, cls=NaturalOrderGroup)
 
 
@@ -518,14 +569,7 @@ def new(path: str = typer.Argument(default=USERCONFIG_FILE,
     """Create a new config file."""
     if os.path.exists(path):
         if check_yes(f'Path "{path}" already exists.\nWould you like to save a copy of it?'):
-            i = 1
-            backup_path = f'{path}.{i}'
-            while os.path.exists(backup_path):
-                i += 1
-                backup_path = f'{path}.{i}'
-            shutil.copy(path, backup_path)
-            print(f'Saved old config to "{backup_path}".')
-            print()
+            config_utils.backup(path)
     ConfigGenerator(path, VALIDATOR).choose_values().save()
 
 

beeflow/common/config_utils.py

@@ -0,0 +1,42 @@
+"""Functions used by the config classes."""
+
+import os
+import shutil
+
+
+def filter_and_validate(config, validator):
+    """Filter and validate the configuration file."""
+    default_keys = list(config['DEFAULT'])
+    config = {sec_name: {key: config[sec_name][key] for key in config[sec_name]
+                         if sec_name == 'DEFAULT' or key not in default_keys} # noqa
+              for sec_name in config}
+    # Validate the config
+    return validator.validate(config)
+
+
+def write_config(file_name, sections):
+    """Write the configuration file."""
+    try:
+        with open(file_name, 'w', encoding='utf-8') as fp:
+            print('# BEE Configuration File', file=fp)
+            for sec_name, section in sections.items():
+                if not section:
+                    continue
+                print(file=fp)
+                print(f'[{sec_name}]', file=fp)
+                for opt_name, value in section.items():
+                    print(f'{opt_name} = {value}', file=fp)
+    except FileNotFoundError:
+        print('Configuration file does not exist!')
+
+
+def backup(fname):
+    """Backup the configuration file."""
+    i = 1
+    backup_path = f'{fname}.{i}'
+    while os.path.exists(backup_path):
+        i += 1
+        backup_path = f'{fname}.{i}'
+    shutil.copy(fname, backup_path)
+    print(f'Saved old config to "{backup_path}".')
+    print()