regscale-cli 6.16.0.0__py3-none-any.whl

regscale/airflow/tasks/cli.py

@@ -0,0 +1,116 @@
+"""Tasks for interacting with the RegScale-CLI."""
+
+from typing import Union
+from uuid import uuid4
+
+from airflow.operators.python import PythonOperator
+
+from regscale.utils.shell import run_command
+
+
+def construct_regscale_login_command(rg_un: str, rg_pw: str) -> str:
+    """Construct the RegScale-CLI login command
+
+    :param str rg_un: a string representing the RegScale User
+    :param str rg_pw: a string representing the RegScale User's password
+    :return: a str of the login command
+    :rtype: str
+    """
+    return f"regscale login --username {rg_un} --password {rg_pw}"
+
+
+def run_login_command(rg_un: str, rg_pw: str) -> bool:
+    """Execute the CLI login command
+
+    :param str rg_un: a string representing the RegScale User
+    :param str rg_pw: a string representing the RegScale User's password
+    :return: a bool indicating success or failure
+    :rtype: bool
+    """
+    return run_command(cmd=construct_regscale_login_command(rg_un=rg_un, rg_pw=rg_pw))
+
+
+def make_cli_task(task_id: str, cmd: Union[str, list], **kwargs: dict) -> PythonOperator:
+    """Construct a cli task
+
+    :param str task_id: a task id to assign to the operator
+    :param Union[str, list] cmd: a str or list of commands to invoke in the cli
+    :param dict **kwargs: additional keyword arguments
+    :return: a PythonOperator task configured to execute run_command
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=task_id,
+        python_callable=run_command,
+        op_kwargs={"cmd": cmd},
+        **kwargs,
+    )
+
+
+def make_login_task(**kwargs: dict) -> PythonOperator:
+    """Construct the regscale login task
+    This login task uses the dag_run.conf keys
+       * `regscale-username`
+       * `regscale-password`
+
+    :param dict **kwargs: additional keyword arguments
+    :return: a PythonOperator task configured to execute run_login_command
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=f"regscale_login_{str(uuid4())[:8]}",
+        python_callable=run_command,
+        op_kwargs={
+            "cmd": 'regscale login --username {{ dag_run.conf["regscale-username"] }}'
+            ' --password {{ dag_run.conf["regscale-password"] }}'
+        },
+        **kwargs,
+    )
+
+
+def make_config_task(key: str, value: str, **kwargs: dict) -> PythonOperator:
+    """Construct a config task to set key to value
+
+    :param str key: str representing the key to set in the config
+    :param str value: str representing the value to set for the key in the config
+    :param dict **kwargs: additional keyword arguments
+    :return: a PythonOperator task configured to execute the run_command to set a value in config
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=f"regscale_config_{key}_{str(uuid4())[:8]}",
+        python_callable=run_command,
+        op_kwargs={"cmd": f"regscale config --param {key} --val {value}"},
+        **kwargs,
+    )
+
+
+def make_set_domain_task(**kwargs: dict) -> PythonOperator:
+    """Construct as config to set the domain, this uses the dag_run.conf["domain"] key
+
+    :param dict **kwargs: additional keyword arguments
+    :return: a PythonOperator that sets the domain
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=f"set_domain_{str(uuid4())[:8]}",
+        python_callable=run_command,
+        op_kwargs={"cmd": 'regscale config --param domain --val {{ dag_run.conf["domain"] }}'},
+        **kwargs,
+    )
+
+
+def make_azure_config_task(**context: dict) -> None:
+    """Construct azure config tasks
+
+    :param dict **context: an airflow context
+    :rtype: None
+    """
+    keys = dict(
+        domain=context["dag_run"].conf.get("domain"),
+        azure365ClientId=context["dag_run"].conf.get("azure365ClientId"),
+        azure365Secret=context["dag_run"].conf.get("azure365Secret"),
+        azure365TenantId=context["dag_run"].conf.get("azure365TenantId"),
+    )
+    for key, value in keys.items():
+        run_command(cmd=f"regscale config --param {key} --val {value}")

regscale/airflow/tasks/click.py

@@ -0,0 +1,73 @@
+"""
+This module contains functions to help Airflow execute a click command.
+"""
+
+import logging
+
+from pydantic import ValidationError
+
+from regscale.models.click_models import ClickCommand
+
+
+def execute_click_command_old(command: ClickCommand, **context: dict) -> None:
+    """Execute a click command
+
+    :param ClickCommand command: a click command to execute
+    :param dict **context: get dag_run from conf
+    :raises ValueError: If required parameters are missing
+    :rtype: None
+    """
+    dag_run = context["dag_run"]
+    try:
+        command_model = command.model_validate(dag_run.conf)
+    except ValidationError as exc:
+        # This will raise a ValidationError if any required parameters are missing
+        raise ValueError(f"Invalid parameters for command '{command.name}': {str(exc)}") from exc
+    valid_params = {name: value for name, value in command_model.dict().items() if name in command.params}
+    command.callback(**valid_params)
+
+
+def execute_click_command(command: ClickCommand, **context: dict) -> None:
+    """Execute a click command
+
+    :param ClickCommand command: a click command to execute
+    :param dict **context: get dag_run from conf
+    :raises ValueError: if no dag_run in context
+    :rtype: None
+    """
+    # Get the dag_run from the context, or raise an error if it's not present
+    dag_run = context.get("dag_run")
+    parameters = command.parameters
+    logging.debug(f"{parameters=}")
+    op_kwargs = {param: context.get(param) for param in parameters}
+    logging.debug(f"{op_kwargs.keys()}")
+    if dag_run is None:
+        raise ValueError("No dag_run in context")
+    if dag_run.conf is None:
+        dag_run.conf = {}
+    # merge the dictionaries, giving precedence to op_kwargs
+    command_defaults = command.defaults
+    command_parameters = {**command_defaults, **dag_run.conf, **op_kwargs}
+    for key, value in command_parameters.items():
+        logging.debug(f"{key=}, {value=}")
+    try:
+        valid_params = {key: value for key, value in command_parameters.items() if key in command.parameters}
+        logging.debug(f"{valid_params=}")
+    except Exception as exc:
+        # Catch any other exceptions that might be raised
+        logging.error(f"Error parsing command {command.name} parameters: {str(exc)}")
+        raise
+    try:
+        if dag_run.conf:
+            from regscale.core.app.application import Application
+
+            _ = Application(config=dag_run.conf)
+            command.context_settings["CONFIG"] = dag_run.conf
+        command.call(**valid_params)
+    except UnboundLocalError as exc:
+        if "callback_results" in str(exc):
+            return None
+        # Catch any exceptions that might be raised by the command callback
+        logging.error(f"Error executing command '{command.name}': {str(exc)}")
+        raise
+    return None

regscale/airflow/tasks/debugging.py

@@ -0,0 +1,9 @@
+"""Tasks for debugging purposes."""
+
+import os
+import logging
+
+
+def set_logging():
+    os.environ["LOGLEVEL"] = "DEBUG"
+    logging.getLogger("airflow.task").setLevel(logging.DEBUG)

regscale/airflow/tasks/groups.py

@@ -0,0 +1,116 @@
+"""Define pre-made TaskGroups for usage across DAGs."""
+
+from uuid import uuid4
+
+from airflow import AirflowException, DAG
+from airflow.models import TaskInstance
+from airflow.operators.python import PythonOperator
+from airflow.utils.db import provide_session
+from airflow.utils.state import State
+from airflow.utils.task_group import TaskGroup
+from airflow.utils.trigger_rule import TriggerRule
+
+from regscale.airflow.hierarchy import AIRFLOW_CLICK_OPERATORS as OPERATORS
+from regscale.airflow.tasks.click import execute_click_command
+
+
+def generate_name(name: str, tag: str = None) -> str:
+    """Generate a unique name for a TaskGroup
+
+    :param str name: the name of the TaskGroup
+    :param str tag: a unique identifier for the TaskGroup
+    :return: a unique name for the TaskGroup
+    :rtype: str
+    """
+    if not tag:
+        tag = str(uuid4())[:8]  # give the task group a unique name for tracking
+    return f"{name}-{tag}"
+
+
+def setup_task_group(
+    dag: DAG,
+    setup_tag: str = None,
+) -> TaskGroup:
+    """Create a TaskGroup for setting up the init.yaml and initialization of the DAG
+
+    :param DAG dag: an Airflow DAG
+    :param str setup_tag: a unique identifier for the task
+    :return: a setup TaskGroup
+    :rtype: TaskGroup
+    """
+    setup_name = generate_name("setup", setup_tag)
+    with TaskGroup(setup_name, dag=dag) as setup:
+        login = PythonOperator(
+            task_id=f"login-{setup_tag}",
+            task_group=setup,
+            python_callable=execute_click_command,
+            op_kwargs={
+                "command": OPERATORS["login"]["command"],
+                "token": '{{ dag_run.conf.get("token") }}',
+                "domain": '{{ dag_run.conf.get("domain") }}',
+            },
+        )
+        login
+        return setup
+
+
+@provide_session
+def _email_on_fail(task, execution_date, dag, session=None, **kwargs):
+    """
+    Send an email if any of the upstream DAGs failed
+
+    :param DAG dag: The DAG to add the operator to
+    :param execution_date: The execution date of the DAG
+    :param session: DAG session
+    :return: The PythonOperator to send an email if any of the DAG tasks fail
+    :rtype: TaskGroup
+    """
+    from regscale.core.app.application import Application
+    from regscale.models import Email
+
+    upstream_task_instances = (
+        session.query(TaskInstance)
+        .filter(
+            TaskInstance.dag_id == dag.dag_id,
+            TaskInstance.execution_date == execution_date,
+            TaskInstance.task_id.in_(task.upstream_task_ids),
+        )
+        .all()
+    )
+    upstream_states = [ti.state for ti in upstream_task_instances]
+    fail_this_task = State.FAILED in upstream_states
+    dag_config = kwargs["dag_run"].conf
+    app = Application(config=dag_config)
+
+    if email := kwargs["dag_run"].conf.get("email"):
+        email = Email(
+            fromEmail="Support@RegScale.com",
+            emailSenderId=app.config["userId"],
+            to=email,
+            subject=f"An Automated Job Has Failed: {dag_config['jobName']} ({dag_config['cadence']})",
+            body=f"{dag.dag_id} with ID: {dag_config['dag_run_id']} job has failed. "
+            + f"Please check the logs for more information: {dag_config['job_url']}",
+        )
+
+    if fail_this_task:
+        if email.send():
+            print(f"Email sent to {email} about {dag.dag_id} job failure.")
+        raise AirflowException("Failing task because one or more upstream tasks failed.")
+
+
+def email_on_fail_operator(dag, email_tag: str = None) -> PythonOperator:
+    """
+    Create a PythonOperator to send an email if any of the DAGs fail
+
+    :param DAG dag: The DAG to add the operator to
+    :param str email_tag: A unique identifier for the task, if not provided one will be generated
+    :return: The PythonOperator to send an email if any of the DAG tasks fail
+    :rtype: TaskGroup
+    """
+    return PythonOperator(
+        task_id=generate_name("email", email_tag),
+        python_callable=_email_on_fail,
+        trigger_rule=TriggerRule.ALL_DONE,
+        provide_context=True,
+        dag=dag,
+    )

regscale/airflow/tasks/init.py

@@ -0,0 +1,60 @@
+"""Initialize init.yaml and set shared keys as variables for the DAG."""
+
+import logging
+from pathlib import Path
+from typing import Optional, Union
+
+import yaml
+from airflow.exceptions import AirflowException
+
+from regscale.airflow.hierarchy import AIRFLOW_CLICK_OPERATORS as OPERATORS
+from regscale.airflow.tasks.click import execute_click_command
+
+
+def get_shared_keys(
+    yaml_file_path: Optional[Union[str, Path]] = Path("init.yaml"),
+    **context,
+) -> list:
+    """Get shared keys between init.yaml and a dag_run_conf obj
+
+    :param Optional[Union[str, Path]] yaml_file_path: the Path to where the yaml file is expected, defaults to Path("init.yaml")
+    :param context: context from Airflow DAG
+    :return: a list of shared keys
+    :rtype: list
+    """
+    if "dag_run" not in context:
+        logging.error(f"context contains {list(context.keys())}")
+    if isinstance(yaml_file_path, str):
+        yaml_file_path = Path(yaml_file_path)
+    yaml_keys = list(yaml.safe_load(yaml_file_path.open("r")).keys())
+    dag_run_conf = context["dag_run"].conf
+    if shared_keys := list(set(list(dag_run_conf.keys())).intersection(set(yaml_keys))):
+        for key in shared_keys:
+            logging.info(f"Shared key: {key}")
+            value = dag_run_conf[key]
+            temp_context = context | {"param": key, "val": value}
+            execute_click_command(command=OPERATORS["config"]["command"], **temp_context)
+    return shared_keys
+
+
+def set_shared_config_values(shared_keys_task: str = None, **context: dict) -> None:
+    """Get the shared keys and set them as a variable
+
+    :param str shared_keys_task: the task_id of the task that will get the shared keys, defaults to None
+    :param dict **context: context from Airflow DAG
+    :raises AirflowException: if no dag_run or dag.conf found or setup_tag specified
+    :rtype: None
+    """
+    logging.info(f"Initial shared_keys_task: {shared_keys_task=}")
+    dag_run_conf = context["dag_run"].conf
+    if not shared_keys_task:
+        if "op_kwargs" not in dag_run_conf:
+            raise AirflowException("No dag_run or dag.conf found or setup_tag specified")
+        shared_keys_task = dag_run_conf["op_kwargs"]["shared_keys_task"]
+    if shared_keys := context["ti"].xcom_pull(key="shared_keys", task_ids=shared_keys_task):
+        for key in shared_keys:
+            value = dag_run_conf[key]
+            temp_context = context | {"param": key, "val": value}
+            execute_click_command(command=OPERATORS["config"]["command"], **temp_context)
+    else:
+        logging.warning(f"No shared keys found: {shared_keys=}")

regscale/airflow/tasks/states.py

@@ -0,0 +1,47 @@
+"""
+This module contains functions to change the state of a DAG run or a task instance.
+"""
+
+from typing import Any
+
+from airflow.exceptions import AirflowException
+from airflow.models import DagRun, TaskInstance
+from airflow.utils.state import State
+
+from airflow import settings
+
+
+def mark_dag_run_as_success(dag_id: Any, execution_date: Any):
+    """
+    This function marks a DAG run as successful.
+
+    :param Any dag_id: the id of the dag
+    :param Any execution_date: the execution date of the dag run
+    """
+    session = settings.Session()
+
+    dag_run = (
+        session.query(DagRun)
+        .filter(
+            DagRun.dag_id == dag_id,
+            DagRun.execution_date == execution_date,
+        )
+        .one_or_none()
+    )
+
+    if dag_run:
+        dag_run.state = State.SUCCESS
+
+        # Update state of the tasks of this dag run
+        session.query(TaskInstance).filter(
+            TaskInstance.dag_id == dag_id,
+            TaskInstance.execution_date == execution_date,
+        ).update({TaskInstance.state: State.SUCCESS})
+
+    session.commit()
+    session.close()
+
+
+def fail_dag_run():
+    """Fail a DAG run on purpose."""
+    raise AirflowException("Failed DAGRun on Purpose")

regscale/airflow/tasks/workflows.py

@@ -0,0 +1,36 @@
+"""Tasks for the workflows to do the following"""
+
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+
+from regscale.airflow.tasks.states import mark_dag_run_as_success, fail_dag_run
+
+
+def build_complete(uid: str, dag: DAG) -> PythonOperator:
+    """Build a complete task
+
+    :param str uid: a unique identifier for the task
+    :param DAG dag: the DAG to add the task to
+    :return: a complete task
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=uid,
+        python_callable=mark_dag_run_as_success,
+        dag=dag,
+    )
+
+
+def build_rejected(uid: str, dag: DAG) -> PythonOperator:
+    """Reject the workflow
+
+    :param str uid: a unique identifier for the task
+    :param DAG dag: the DAG to add the task to
+    :return: a rejected task
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=uid,
+        python_callable=fail_dag_run,
+        dag=dag,
+    )

regscale/ansible/__init__.py

@@ -0,0 +1,9 @@
+"""Provide ansible integraitons."""
+
+try:
+    import ansible
+except ImportError:
+    raise ImportError(
+        "The `ansible` package is required for using the regscale.ansible package. "
+        "To install it, run: `pip install regscale-cli[ansible]`"
+    )

regscale/core/app/__init__.py

@@ -0,0 +1,3 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from regscale.core.app.logz import create_logger