Implement changes from generic module

deployability/launchers/workflow_engine.py

@@ -27,22 +27,11 @@ def parse_arguments() -> argparse.Namespace:
     parser.add_argument('--schema_file', required=False, type=str, help='Path to the schema file (YAML format)')
     return parser.parse_args()
 
-
-def setup_logger(log_level: str) -> None:
-    """Setup logger."""
-    logger = logging.getLogger()
-    console_handler = colorlog.StreamHandler()
-    console_handler.setFormatter(colorlog.ColoredFormatter("%(log_color)s[%(asctime)s] [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
-    logger.addHandler(console_handler)
-    logger.setLevel(log_level)
-
-
 def main() -> None:
     """Main entry point."""
 
     args = parse_arguments()
-    setup_logger(args.log_level)
-    processor = WorkflowProcessor(**dict(InputPayload(**vars(parse_arguments()))))
+    processor = WorkflowProcessor(**dict(InputPayload(**vars(args))))
     processor.run()
 
 

deployability/modules/__init__.py

@@ -1,4 +1,3 @@
 from .provision import Provision
 from .generic import Ansible
 from .allocation import Allocator
-from .generic import SchemaValidator

deployability/modules/allocation/allocation.py

@@ -2,8 +2,9 @@
 
 from pathlib import Path
 
-from .generic import Instance, Provider, models
 from .aws.provider import AWSProvider, AWSConfig
+from .generic import Instance, Provider, models
+from .generic.utils import logger
 from .vagrant.provider import VagrantProvider, VagrantConfig
 
 
@@ -26,10 +27,10 @@ def run(cls, payload: models.InputPayload) -> None:
         payload = models.InputPayload(**dict(payload))
         # Detect the action and call the appropriate method.
         if payload.action == 'create':
-            print(f"Creating instance at {payload.working_dir}")
+            logger.info(f"Creating instance at {payload.working_dir}")
             return cls.__create(payload)
         elif payload.action == 'delete':
-            print(f"Deleting instance from trackfile {payload.track_output}")
+            logger.info(f"Deleting instance from trackfile {payload.track_output}")
             return cls.__delete(payload)
 
     # Internal methods
@@ -48,10 +49,10 @@ def __create(cls, payload: models.CreationPayload):
         config = cls.___get_custom_config(payload)
         instance = provider.create_instance(
             payload.working_dir, instance_params, config)
-        print(f"Instance {instance.identifier} created.")
-        # Start the instance
+        logger.info(f"Instance {instance.identifier} created.")
+        # Start the instance    
         instance.start()
-        print(f"Instance {instance.identifier} started.")
+        logger.info(f"Instance {instance.identifier} started.")
         # Generate the inventory and track files.
         cls.__generate_inventory(instance, payload.inventory_output)
         cls.__generate_track_file(instance, payload.provider, payload.track_output)
@@ -71,7 +72,7 @@ def __delete(cls, payload: models.DeletionPayload) -> None:
             track = models.TrackOutput(**yaml.safe_load(f))
         provider = PROVIDERS[track.provider]()
         provider.destroy_instance(track.instance_dir, track.identifier)
-        print(f"Instance {track.identifier} deleted.")
+        logger.info(f"Instance {track.identifier} deleted.")
 
     @staticmethod
     def ___get_custom_config(payload: models.CreationPayload) -> models.ProviderConfig | None:
@@ -85,11 +86,13 @@ def ___get_custom_config(payload: models.CreationPayload) -> models.ProviderConf
         Returns:
             ProviderConfig: The configuration object.
         """
-        if not payload.custom_provider_config:
+        config = payload.custom_provider_config
+        if not config:
             return None
         # Read the custom config file and validate it.
         config_model: models.ProviderConfig = CONFIGS[payload.provider]
-        with open(payload.custom_provider_config, 'r') as f:
+        with open(config, 'r') as f:
+            logger.info(f"Using custom provider config from {config}")
             config = config_model(**yaml.safe_load(f))
         return config
 
@@ -112,7 +115,7 @@ def __generate_inventory(instance: Instance, inventory_path: Path) -> None:
                                            ansible_ssh_private_key_file=str(ssh_config.private_key))
         with open(inventory_path, 'w') as f:
             yaml.dump(inventory.model_dump(), f)
-        print(f"\nInventory file generated at {inventory_path}")
+        logger.info(f"Inventory file generated at {inventory_path}")
 
     @staticmethod
     def __generate_track_file(instance: Instance, provider_name: str,  track_path: Path) -> None:
@@ -133,4 +136,4 @@ def __generate_track_file(instance: Instance, provider_name: str,  track_path: P
                                    key_path=str(instance.credentials.key_path))
         with open(track_path, 'w') as f:
             yaml.dump(track.model_dump(), f)
-        print(f"\nTrack file generated at {track_path}")
+        logger.info(f"Track file generated at {track_path}")

deployability/modules/allocation/aws/credentials.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 from modules.allocation.generic import Credentials
+from modules.allocation.generic.utils import logger
 
 
 class AWSCredentials(Credentials):
@@ -46,6 +47,7 @@ def generate(self, base_dir: str | Path, name: str, overwrite: bool = False) ->
 
         # Validate base directory
         if not base_dir.exists():
+            logger.debug(f"Creating base directory: {base_dir}")
             base_dir.mkdir(parents=True, exist_ok=True)
         elif base_dir.is_file():
             raise self.CredentialsError(f"Invalid base directory: {base_dir}")
@@ -57,6 +59,7 @@ def generate(self, base_dir: str | Path, name: str, overwrite: bool = False) ->
                 if not overwrite:
                     raise self.CredentialsError(f"Key pair {name} already exists.")
                 else:
+                    logger.warning(f"Key pair {name} already exists. Overwriting.")
                     key_pair.delete()
         except ClientError:
             pass
@@ -112,6 +115,7 @@ def load(self, name: str) -> str:
     def delete(self) -> None:
         """Deletes the key pair."""
         if not self.name:
+            logger.warning(f"Key pair doesn't exist. Skipping deletion.")
             return
 
         try:
@@ -122,6 +126,7 @@ def delete(self) -> None:
 
         # Remove the local private key file
         if self.key_path:
+            logger.debug(f"Deleting private key: {self.key_path}")
             Path(self.key_path).unlink()
 
         # Clear instance attributes

deployability/modules/allocation/aws/instance.py

@@ -1,7 +1,9 @@
 import boto3
+
 from pathlib import Path
 from modules.allocation.generic import Instance
 from modules.allocation.generic.models import ConnectionInfo
+from modules.allocation.generic.utils import logger
 from .credentials import AWSCredentials
 
 
@@ -32,6 +34,7 @@ def __init__(self, path: str | Path, identifier: str, credentials: AWSCredential
         self._client = boto3.resource('ec2')
         self._instance = self._client.Instance(self.identifier)
         if not self.credentials:
+            logger.debug(f"No credentials found. Loading from instance directory.")
             self.credentials = self.__get_credentials()
 
     def start(self) -> None:

deployability/modules/allocation/aws/provider.py

@@ -5,6 +5,7 @@
 
 from modules.allocation.generic import Provider
 from modules.allocation.generic.models import CreationPayload
+from modules.allocation.generic.utils import logger
 from .credentials import AWSCredentials
 from .instance import AWSInstance
 from .models import AWSConfig
@@ -38,21 +39,25 @@ def _create_instance(cls, base_dir: Path, params: CreationPayload, config: AWSCo
         temp_dir = base_dir / temp_id
         credentials = AWSCredentials()
         if not config:
+            logger.debug(f"No config provided. Generating from payload")
             # Generate the credentials.
             credentials.generate(temp_dir, temp_id.split('-')[-1] + '_key')
             # Parse the config if it is not provided.
             config = cls.__parse_config(params, credentials)
         else:
+            logger.debug(f"Using provided config")
             # Load the existing credentials.
             credentials.load(config.key_name)
             # Create the temp directory. 
             # TODO: Review this on the credentials refactor.
             if not temp_dir.exists():
+                logger.debug(f"Creating temp directory: {temp_dir}")
                 temp_dir.mkdir(parents=True, exist_ok=True)
         # Generate the instance.
         instance_id = cls.__create_ec2_instance(config)
         # Rename the temp directory to its real name.
         instance_dir = Path(base_dir, instance_id)
+        logger.debug(f"Renaming temp {temp_dir} directory to {instance_dir}")
         os.rename(temp_dir, instance_dir)
         credentials.key_path = (instance_dir / credentials.name).with_suffix('.pem')
 
@@ -83,6 +88,7 @@ def _destroy_instance(cls, instance_dir: str, identifier: str) -> None:
         """
         instance = AWSInstance(instance_dir, identifier)
         if instance.credentials:
+            logger.debug(f"Deleting credentials: {instance.credentials.key_path}")
             instance.credentials.delete()
         instance.delete()
 

deployability/modules/allocation/generic/provider.py

@@ -7,6 +7,7 @@
 
 from .instance import Instance
 from .models import CreationPayload, ProviderConfig
+from .utils import logger
 
 
 class Provider(ABC):
@@ -76,11 +77,11 @@ def load_instance(cls, instance_dir: str | Path, instance_id: str) -> Instance:
             Instance: The loaded instance.
 
         Raises:
-            Exception: If the instance directory does not exist.
+            ValueError: If the instance directory does not exist.
         """
         instance_dir = Path(instance_dir)
         if not instance_dir.exists():
-            raise Exception(f"Instance path {instance_dir} does not exist")
+            raise ValueError(f"Instance path {instance_dir} does not exist")
         return cls._load_instance(instance_dir, instance_id)
 
     @classmethod

deployability/modules/allocation/generic/utils.py

@@ -0,0 +1,3 @@
+from modules.generic.logger import Logger
+
+logger = Logger("allocator").get_logger()

deployability/modules/allocation/vagrant/credentials.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 from modules.allocation.generic import Credentials
+from modules.allocation.generic.utils import logger
 
 
 class VagrantCredentials(Credentials):
@@ -36,9 +37,12 @@ def generate(self, base_dir: str | Path, name: str) -> Path:
             CredentialsError: This exception is raised if there's an error during the key creation process.
         """
         if self.key_path and self.key_id:
+            logger.warning(f"Key pair already exists: {self.key_path}")
             return self.key_path
+
         base_dir = Path(base_dir)
         if not base_dir.exists():
+            logger.debug(f"Creating base directory: {base_dir}")
             base_dir.mkdir(parents=True, exist_ok=True)
         elif Path(base_dir).is_file():
             raise self.CredentialsError(f"Invalid base directory: {base_dir}")
@@ -47,6 +51,7 @@ def generate(self, base_dir: str | Path, name: str) -> Path:
         public_key_path = private_key_path.with_suffix(".pub")
         # Delete the existing key pair if it exists.
         if private_key_path.exists():
+            logger.warning(f"Key pair already exists: {private_key_path}")
             return self.load(base_dir, name)
         elif private_key_path.exists():
             private_key_path.unlink()
@@ -63,8 +68,8 @@ def generate(self, base_dir: str | Path, name: str) -> Path:
                                 capture_output=True, text=True)
         os.chmod(private_key_path, 0o600)
         if output.returncode != 0:
-            raise self.CredentialsError(
-                f"Error creating key pair: {output.stderr}")
+            raise self.CredentialsError(f"Error creating key pair: {output.stderr}")
+
         # Save instance attributes.
         self.name = name
         self.key_id = name
@@ -83,7 +88,7 @@ def load(self, path: str | Path) -> None:
         """
         key_path = Path(path)
         if not key_path.exists() or not key_path.is_file():
-            raise self.CredentialsError(f"Invalid path {key_path}.")
+            raise self.CredentialsError(f"Invalid key path {key_path}.")
         self.key_path = key_path
         self.name = key_path.name
         self.key_id = key_path.name
@@ -93,6 +98,8 @@ def delete(self) -> None:
         Deletes the key pair from the file system.
         """
         if not self.key_path.exists():
+            logger.warning(f"Key pair doesn't exist: {self.key_path}.\
+                            Skipping deletion.")
             return
         Path(self.key_path).unlink()
         Path(self.key_path.with_suffix(".pub")).unlink()

deployability/modules/allocation/vagrant/instance.py

@@ -5,6 +5,7 @@
 
 from modules.allocation.generic import Instance
 from modules.allocation.generic.models import ConnectionInfo
+from modules.allocation.generic.utils import logger
 from .credentials import VagrantCredentials
 
 
@@ -65,6 +66,8 @@ def delete(self) -> None:
             None
         """
         if "not created" in self.status():
+            logger.warning(f"Instance {self.identifier} is not created.\
+                            Skipping deletion.")
             return
         self.__run_vagrant_command(['destroy', '-f'])
 
@@ -86,6 +89,8 @@ def ssh_connection_info(self) -> ConnectionInfo:
             ConnectionInfo: The SSH configuration of the VM.
         """
         if not 'running' in self.status():
+            logger.debug(f"Instance {self.identifier} is not running.\
+                          Starting it.")
             self.start()
         output = self.__run_vagrant_command('ssh-config')
         patterns = {'hostname': r'HostName (.*)',
@@ -99,8 +104,10 @@ def ssh_connection_info(self) -> ConnectionInfo:
             if match:
                 ssh_config[key] = str(match.group(1)).strip("\r")
             else:
-                raise ValueError(f"Couldn't find {key} in vagrant ssh-config")
+                logger.error(f"Couldn't find {key} in ssh-config")
+                return None
         if self.credentials:
+            logger.debug(f"Using provided credentials")
             ssh_config['private_key'] = str(self.credentials.key_path)
         return ConnectionInfo(**ssh_config)
 
@@ -124,11 +131,11 @@ def __run_vagrant_command(self, command: str | list) -> str:
                                     stderr=subprocess.PIPE)
 
             if stderr := output.stderr.decode("utf-8"):
-                print(stderr)
-                print(output.stdout.decode("utf-8"))
+                logger.error(f"Command failed: {stderr}")
+                return None
             return output.stdout.decode("utf-8")
         except subprocess.CalledProcessError as e:
-            print(e)
+            logger.error(f"Command failed: {e.stderr}")
             return None
 
     def __parse_vagrant_status(self, message: str) -> str:

deployability/modules/allocation/vagrant/provider.py

@@ -6,6 +6,7 @@
 
 from modules.allocation.generic import Provider
 from modules.allocation.generic.models import CreationPayload
+from modules.allocation.generic.utils import logger
 from .credentials import VagrantCredentials
 from .instance import VagrantInstance
 from .models import VagrantConfig
@@ -40,14 +41,17 @@ def _create_instance(cls, base_dir: Path, params: CreationPayload, config: Vagra
         instance_dir.mkdir(parents=True, exist_ok=True)
         credentials = VagrantCredentials()
         if not config:
+            logger.debug(f"No config provided. Generating from payload")
             # Generate the credentials.
             credentials.generate(instance_dir, 'instance_key')
             # Parse the config if it is not provided.
             config = cls.__parse_config(params, credentials)
         else:
+            logger.debug(f"Using provided config")
             credentials.load(config.public_key)
         # Create the Vagrantfile.
         cls.__create_vagrantfile(instance_dir, config)
+        logger.debug(f"Vagrantfile created. Creating instance.")
         return VagrantInstance(instance_dir, instance_id, credentials)
 
     @staticmethod
@@ -64,8 +68,8 @@ def _load_instance(instance_dir: Path, identifier: str) -> VagrantInstance:
         """
         return VagrantInstance(instance_dir, identifier)
 
-    @staticmethod
-    def _destroy_instance(instance_dir: Path, identifier: str) -> None:
+    @classmethod
+    def _destroy_instance(cls, instance_dir: Path, identifier: str) -> None:
         """
         Destroys a Vagrant instance.
 
@@ -77,6 +81,7 @@ def _destroy_instance(instance_dir: Path, identifier: str) -> None:
             None
         """
         instance = VagrantInstance(instance_dir, identifier)
+        logger.debug(f"Destroying instance {identifier}")
         instance.delete()
 
     @classmethod

deployability/modules/generic/__init__.py

@@ -1,2 +1 @@
 from .ansible import Ansible, Inventory
-from .schemaValidator import SchemaValidator