regscale-cli 6.16.0.0__py3-none-any.whl

regscale/__init__.py

@@ -0,0 +1 @@
+__version__ = "6.16.0.0"

regscale/airflow/__init__.py

@@ -0,0 +1,9 @@
+"""Provide Airflow integrations with RegScale."""
+
+try:
+    import airflow
+except ImportError:
+    raise ImportError(
+        "The 'apache-airflow' package is required for using regscale.airflow. "
+        "To install it, run: `pip install regscale-cli[airflow]`"
+    )

regscale/airflow/azure/__init__.py

@@ -0,0 +1,9 @@
+"""Prevent usage and imports that will break when used outside of airflow-azure extra."""
+
+import sys
+
+try:
+    from azure.storage.blob import BlobServiceClient
+except ImportError:
+    print("To use Azure Blob Storage features, you need to install the [airflow-azure] extra package.")
+    sys.exit(1)

regscale/airflow/azure/cli.py

@@ -0,0 +1,89 @@
+"""Provide a CLI for uploading DAGs to Azure Blob Storage."""
+
+import os
+import sys
+import click
+
+from regscale.airflow.azure.upload_dags import (
+    upload_dag_to_blob_storage,
+    upload_dags_to_blob_storage,
+)
+
+
+@click.group(name="dags")
+def cli():
+    """Upload DAGs or files to Azure Blob Storage."""
+    pass
+
+
+@cli.command()
+@click.option(
+    "--file",
+    "-f",
+    "file_path",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to the DAG file to upload.",
+)
+@click.option(
+    "--conn-string",
+    "-c",
+    "connection_string",
+    default=os.getenv("AZURE_STORAGE_CONNECTION_STRING", None),
+    help="Azure Blob Storage connection string.",
+)
+@click.option(
+    "--container",
+    "-n",
+    "container_name",
+    default=os.getenv("AZURE_STORAGE_CONTAINER_NAME", "dags"),
+    help="Azure Blob Storage container name.",
+)
+def upload_dag(file_path: str, connection_string: str, container_name: str) -> None:
+    """Upload a single DAG to Azure Blob Storage."""
+    if connection_string is None:
+        click.echo("You can set the connection string with the AZURE_STORAGE_CONNECTION_STRING environment variable.")
+    if container_name is None:
+        click.echo("You can set the container name with the AZURE_STORAGE_CONTAINER_NAME environment variable.")
+    if connection_string is None or container_name is None:
+        click.echo("Please provide the connection string and container name for Azure Blob Storage.")
+        sys.exit(1)
+    upload_dag_to_blob_storage(
+        connection_string=connection_string,
+        container_name=container_name,
+        blob_name=file_path,
+        file_path=file_path,
+    )
+
+
+@cli.command()
+@click.option("--path", "-p", "path", default="airflow/dags/", help="Path to the DAGs folder.")
+@click.option(
+    "--conn-string",
+    "-c",
+    "connection_string",
+    default=os.getenv("AZURE_STORAGE_CONNECTION_STRING", None),
+    help="Azure Blob Storage connection string.",
+)
+@click.option(
+    "--container",
+    "-n",
+    "container_name",
+    default=os.getenv("AZURE_STORAGE_CONTAINER_NAME", "dags"),
+    help="Azure Blob Storage container name.",
+)
+def upload_dags(path: str, connection_string: str, container_name: str) -> None:
+    """Upload DAGs to Azure Blob Storage."""
+    if connection_string is None or container_name is None:
+        click.echo("Please provide the connection string and container name for Azure Blob Storage.")
+    if connection_string is None:
+        click.echo("You can set the connection string with the AZURE_STORAGE_CONNECTION_STRING environment variable.")
+    if container_name is None:
+        click.echo("You can set the container name with the AZURE_STORAGE_CONTAINER_NAME environment variable.")
+    if connection_string is None or container_name is None:
+        sys.exit(1)
+    upload_dags_to_blob_storage(
+        path=path,
+        connection_string=connection_string,
+        container_name=container_name,
+    )

regscale/airflow/azure/upload_dags.py

@@ -0,0 +1,116 @@
+"""Provide functions and an entrypoint for uploading DAGs to Azure Blob Storage."""
+
+import os
+import sys
+from pathlib import Path
+from typing import Optional, Union, List, Tuple
+
+from azure.storage.blob import BlobServiceClient
+
+from regscale.core.app.logz import create_logger
+
+
+def upload_dag_to_blob_storage(
+    connection_string: str,
+    container_name: str,
+    blob_name: str,
+    file_path: Union[str, Path],
+    exit_on_failure: bool = True,
+) -> None:
+    """Upload a DAG to Azure Blob Storage
+
+    :param str connection_string: Azure Blob Storage connection string
+    :param str container_name: Azure Blob Storage container name
+    :param str blob_name: Azure Blob Storage blob name
+    :param Union[str, Path] file_path: Path to the DAG file
+    :param bool exit_on_failure: Whether to exit on failure, default True
+    :rtype: None
+    """
+    logger = create_logger()
+    if isinstance(file_path, str):
+        file_path = Path(file_path)
+    try:
+        blob_service_client = BlobServiceClient.from_connection_string(connection_string)
+        blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name)
+        blob_client.upload_blob(file_path.read_bytes(), overwrite=True)
+        logger.info(f"{file_path} with {blob_name} uploaded to Azure Blob Storage")
+    except Exception as e:
+        logger.error(f"Failed to upload DAG to Azure Blob Storage: {e}")
+        if exit_on_failure:
+            sys.exit(1)
+
+
+def retrieve_dags_to_upload(path: Union[str, Path], file_extension: str = ".py") -> List[Tuple[Path, str]]:
+    """
+    Retrieve DAGs to upload to Azure Blob Storage
+
+    :param Union[str, Path] path: Path to the DAGs folder
+    :param str file_extension: File extension of the DAGs, default .py
+    :return: List of DAGs to upload
+    :rtype: List[Tuple[Path, str]]
+    """
+    if isinstance(path, str):
+        path = Path(path)
+    dags = path.glob(f"*{file_extension}")
+    return [(dag, dag.name) for dag in dags]
+
+
+def upload_dags_to_blob_storage(
+    path: str = "airflow/dags/",
+    connection_string: Optional[str] = None,
+    container_name: Optional[str] = None,
+    exit_on_failure: bool = True,
+) -> None:
+    """Upload DAGs to Azure Blob Storage
+
+    :param str path: Path to the DAGs folder, default airflow/dags/
+    :param Optional[str] connection_string: Azure Blob Storage connection string
+    :param Optional[str] container_name: Azure Blob Storage container name
+    :param bool exit_on_failure: Whether to exit on failure, default True
+    :rtype: None
+    """
+    logger = create_logger()
+    if connection_string is None:
+        connection_string = os.environ.get("AZURE_STORAGE_CONNECTION_STRING")
+    if container_name is None:
+        container_name = os.environ.get("AZURE_STORAGE_CONTAINER_NAME")
+    if connection_string is None or container_name is None:
+        logger.error("Please provide the connection string and container name for Azure Blob Storage.")
+        if exit_on_failure:
+            sys.exit(1)
+    dags = retrieve_dags_to_upload(path)
+    for dag, blob_name in dags:
+        upload_dag_to_blob_storage(
+            connection_string=connection_string,
+            container_name=container_name,
+            blob_name=blob_name,
+            file_path=dag,
+            exit_on_failure=exit_on_failure,
+        )
+
+
+def main():
+    """Entrypoint for uploading DAGs to Azure Blob Storage."""
+    logger = create_logger()
+    if len(sys.argv) == 2:
+        connection_string = sys.argv[1]
+        container_name = sys.argv[2]
+    else:
+        connection_string = os.environ.get("ADS_CONN_STRING")
+        container_name = os.environ.get("AZURE_STORAGE_CONTAINER_NAME", "dags")
+    if connection_string is None or container_name is None:
+        logger.error("Please provide the connection string and container name for Azure Blob Storage.")
+        if connection_string is None:
+            logger.error("Connection string not provided.")
+        if container_name is None:
+            logger.error("Container name not provided.")
+        sys.exit(1)
+    upload_dags_to_blob_storage(
+        connection_string=connection_string,
+        container_name=container_name,
+        exit_on_failure=False,
+    )
+
+
+if __name__ == "__main__":
+    main()

regscale/airflow/click_dags.py

@@ -0,0 +1,127 @@
+"""Create functions to generate a DAG based on click hierarchy."""
+
+import os
+from datetime import datetime, timezone
+from typing import Optional
+
+import click
+from airflow.operators.python import PythonOperator
+
+from airflow import DAG
+from regscale.airflow.tasks.cli import make_login_task, make_set_domain_task
+from regscale.airflow.tasks.click import execute_click_command
+from regscale.models.click_models import ClickCommand
+from regscale.regscale import cli
+
+
+def generate_operator_for_command(
+    command: ClickCommand,
+    dag: DAG,
+    command_name: str = None,
+) -> PythonOperator:
+    """Generate an Operator for click command
+
+    :param ClickCommand command: a click.Command instruction
+    :param DAG dag: an Airflow DAG
+    :param str command_name: an optional string for command_name, will default to click.Command.name
+    :return: A PythonOperator configured with the click command.name
+    :rtype: PythonOperator
+    """
+    return PythonOperator(
+        task_id=command_name or command.name,  # TODO - also name based on group id
+        python_callable=execute_click_command,
+        op_kwargs={"command": command},
+        dag=dag,
+    )
+
+
+def generate_dag_for_group(
+    group: click.Group,
+    default_args: dict,
+    command_name: str = None,
+) -> DAG:
+    """Generate a dag for a click group
+
+    :param click.Group group: a click Group object to generate dags for
+    :param dict default_args: dict to be passed when creating the DAGs
+    :param str command_name: an optional string for command_name, will default to click.Command.name
+    :return: a dag for each click group
+    :rtype: DAG
+    """
+    if command_name is None:
+        command_name = group.name
+    dag = DAG(
+        command_name,
+        default_args=default_args,
+        description=f"DAG for Click Group: {group.name}",
+        schedule_interval=None,
+    )
+
+    # iterate through the group value commands
+    for command in group.commands.values():
+        if isinstance(command, click.Group):
+            # if it is click.Group, call this recursively, adding to the command_name
+            generate_dag_for_group(
+                command,
+                default_args,
+                command_name="__".join([command_name, command.name]),
+            )
+        else:
+            login_task = make_login_task(dag=dag)
+            domain_task = make_set_domain_task(dag=dag)
+            operator = generate_operator_for_command(
+                command=command,
+                dag=dag,
+                command_name="__".join([command_name, command.name]),
+            )
+            # assign the relationships for the operators
+            domain_task >> login_task >> operator
+    return dag
+
+
+def generate_dags_for_click_app(
+    app: Optional[click.Group] = None,
+    default_args: Optional[dict] = None,
+    command_name: Optional[str] = None,
+) -> list:
+    """Generate DAGs for a click app
+
+    :param Optional[click.Group] app: a click app to generate dags for, defaults to None
+    :param Optional[dict] default_args: dict to be passed when creating the DAGs, defaults to None
+    :param Optional[str] command_name: the name of the command to generate the DAG for, defaults to None
+    :return: a list of DAGs for the click app
+    :rtype: list
+    """
+    # if app is not passed, use the cli
+    if app is None:
+        app = cli
+    if not default_args:
+        default_args = dict(
+            owner=os.getenv("REGSCALE_AIRFLOW_USER", "regscale"),
+            start_date=datetime.now(timezone.utc),
+        )
+    if not command_name:
+        command_name = "regscale" if app.name == "cli" else app.name
+    dags = []
+    for group in app.commands.values():
+        if isinstance(group, click.Group):
+            command_name = "regscale" if group.name == "cli" else group.name
+            for sub_group in group.commands.values():
+                if isinstance(sub_group, click.Group):
+                    generate_dag_for_group(app=sub_group, default_args=default_args)
+                else:
+                    dag = generate_dag_for_group(
+                        group=group,
+                        default_args=default_args,
+                        command_name=command_name,
+                    )
+                    dags.append(dag)
+        elif isinstance(group, click.Command):
+            dag = generate_dag_for_group(
+                group=group,
+                default_args=default_args,
+                command_name=command_name,
+            )
+            dags.append(dag)
+
+    return dags

regscale/airflow/click_mixins.py

@@ -0,0 +1,82 @@
+"""Provide mixins for click models."""
+
+from typing import Any
+from uuid import uuid4
+from airflow.operators.python import PythonOperator
+
+from regscale.models.click_models import ClickGroup, ClickCommand
+from regscale.airflow.tasks.click import execute_click_command
+
+
+class AirflowOperatorMixin:
+    """Mixin to the ClickGroup to flatten for an Airflow Operator"""
+
+    def flatten_operator(self):
+        """Flatten the group to a dictionary of PythonOperator objects"""
+        operators_ = {}
+
+        def _flatten_operator(group_: Any, prefix: str = "") -> dict:
+            """Flatten the group to a dictionary of PythonOperator objects
+
+            :param Any group_: a click.Group object to generate dags for
+            :param str prefix: a string to prefix the command name with, defaults to ""
+            :return: a dictionary of PythonOperator objects
+            :rtype: dict
+            """
+            for name, cmd in group_.commands.items():
+                if isinstance(cmd, ClickCommand):
+                    cmd_name = cmd.name if prefix == "" else f"{prefix}__{cmd.name}"
+
+                    def _make_operator_wrapper(
+                        cmd_name_: str, cmd_: ClickCommand, suffix: str = None, **kwargs
+                    ) -> PythonOperator:
+                        """Create a wrapper to make a python operator."""
+                        if suffix is None:
+                            suffix = str(uuid4())[:8]
+                        op_kwargs = None
+                        if "op_kwargs" in kwargs:
+                            op_kwargs = kwargs.pop("op_kwargs")
+                        inputs = dict(
+                            task_id=f"{cmd_name_}-{suffix}",
+                            python_callable=execute_click_command,
+                        )
+                        if kwargs:
+                            inputs |= kwargs
+                        if op_kwargs:
+                            inputs |= {"op_kwargs": {"command": cmd_, **op_kwargs}}
+                        else:
+                            inputs |= {"op_kwargs": {"command": cmd_}}
+                        return PythonOperator(**inputs)
+
+                    def _construct_lambda_wrapper(cmd_name_, cmd_, suffix: str = None, **kwargs):
+                        return _make_operator_wrapper(cmd_name_=cmd_name_, cmd_=cmd_, suffix=suffix, **kwargs)
+
+                    operators_[cmd_name] = {
+                        "operator": PythonOperator(
+                            task_id=cmd_name,
+                            python_callable=execute_click_command,
+                            op_kwargs={"command": cmd},
+                        ),
+                        "lambda": lambda cmd_name_=cmd_name, cmd_=cmd, suffix=None, **kwargs: _construct_lambda_wrapper(
+                            cmd_name_=cmd_name_, cmd_=cmd_, suffix=suffix, **kwargs
+                        ),
+                        "command": cmd,
+                    }
+                elif isinstance(cmd, ClickGroup):
+                    new_prefix = f"{prefix}__{cmd.group_name}" if prefix else cmd.group_name
+                    _flatten_operator(cmd, new_prefix)
+
+        _flatten_operator(self)
+        return operators_
+
+
+class AirflowClickGroup(AirflowOperatorMixin, ClickGroup):
+    """Initialize the AirflowClickGroup object."""
+
+
+# here's an example of how to generate the OPERATORS AirflowClickGroup class
+if __name__ == "__main__":
+    from regscale.regscale import cli
+
+    group = AirflowClickGroup.from_group(cli, prefix="regscale")
+    operators = group.flatten_operator()

regscale/airflow/config.py

@@ -0,0 +1,25 @@
+"""Provide configurations for Airflow."""
+
+from datetime import datetime, timedelta
+
+
+def yesterday():
+    """Return yesterday from now in datetime"""
+    return datetime.now() - timedelta(days=1)
+
+
+DEFAULT_ARGS = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    # 'start_date': (datetime.now() - timedelta(days=1)).date(),  # left here to show we intentionally disable
+    "email": ["airflow@regscale.com"],
+    "email_on_failure": False,
+    "email_on_retry": False,
+    "retries": 2,
+    "retry_delay": timedelta(minutes=2),
+    "execution_timeout": timedelta(hours=3),
+    # 'queue': 'whatever queue we want to implement',  # left here for an example
+    # 'pool': 'backfill',  # another example default arg
+    # 'priority_weight': 10,  # give this a high priority weight
+    # 'end_date': datetime(2038, 1, 1),  # left to show that end dates can be set
+}

regscale/airflow/factories/connections.py

@@ -0,0 +1,58 @@
+"""Provide connection management functions."""
+
+import os
+from typing import Literal
+
+from airflow.models import Connection
+from airflow.operators.python import PythonOperator
+
+from airflow import settings
+
+
+def create_connection_operator(
+    conn_id: str,
+    conn_type: Literal["postgres", "mssql"],
+    **kwargs: dict,
+) -> PythonOperator:
+    """Create a connection if it does not exist
+
+    :param str conn_id: the connection id
+    :param Literal["postgres", "mssql"] conn_type: the connection type
+    :param dict **kwargs: additional keyword arguments
+    :return: a PythonOperator to create the connection
+    :rtype: PythonOperator
+    """
+
+    def _create_connection():
+        """Create a connection"""
+        session = settings.Session()
+        # Check if connection exists
+        conn = session.query(Connection).filter(Connection.conn_id == conn_id).first()
+        # If connection does not exist, create it
+        if conn is None:
+            new_conn = Connection(
+                conn_id=conn_id,
+                conn_type=conn_type,
+                host=os.getenv("PLATFORM_DB_HOST") or "atlas",
+                # FIXME - this envar needs added to the instance container app
+                schema=os.getenv("PLATFORM_DB_NAME") or "airflowtest-sqlDatabase",
+                login=os.getenv("PLATFORM_DB_USER"),
+                password=os.getenv("PLATFORM_DB_PASSWORD"),
+                port=int(
+                    os.getenv(
+                        "PLATFORM_DB_PORT",
+                    )
+                    or (5432 if conn_type == "postgres" else 1433)
+                ),
+            )
+            session.add(new_conn)
+            session.commit()
+            print("Connection created.")
+        else:
+            print("Connection already exists.")
+
+    return PythonOperator(
+        task_id=f"create-connection-{conn_id}",
+        python_callable=_create_connection,
+        **kwargs,
+    )

regscale/airflow/factories/workflows.py

@@ -0,0 +1,78 @@
+"""Provide workflow factory functions."""
+
+from typing import Union
+from uuid import uuid4
+
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+from airflow.utils.task_group import TaskGroup
+
+from regscale.airflow.sensors.sql import build_sql_sensor_xcon
+from regscale.airflow.sessions.sql.sql_server_queries import (
+    CHECK_IF_COMPLETED_SQL_QUERY,
+)
+from regscale.airflow.tasks.branches import tri_branch_func
+from regscale.airflow.tasks.workflows import (
+    build_complete,
+    build_rejected,
+)
+
+
+def workflow_listener_factory(
+    workflow_name: str,
+    dag: DAG,
+    next_step_name: str,
+    next_step: Union[TaskGroup, PythonOperator],
+    unique: bool = False,
+    **kwargs,
+) -> TaskGroup:
+    """Build a workflow listener task group
+
+    :param str workflow_name: the name of the workflow
+    :param DAG dag: the DAG to add the task group to
+    :param str next_step_name: the name of the next step
+    :param Union[TaskGroup, PythonOperator] next_step: the next step
+    :param dict kwargs: keyword arguments to pass to the task group
+    :param bool unique: whether to make the task group unique
+    :return: a workflow listener task group
+    :rtype: TaskGroup
+    """
+    uid = None
+    if unique:
+        uid = str(uuid4())[:8]
+        name = f"{workflow_name}-{uid}"
+    else:
+        name = f"{workflow_name}"
+    with TaskGroup(
+        group_id=name,
+        dag=dag,
+        **kwargs,
+    ) as task_group:
+        # You need to have this factory upstream from the is_completed_listener_name
+        is_completed_listener_name = f"{name}-listener"
+        complete_name = f"{name}-complete"
+        # this task needs to be the name of the next listener
+        rejected_name = f"{name}-rejected"
+        is_completed = build_sql_sensor_xcon(
+            sql=CHECK_IF_COMPLETED_SQL_QUERY,
+            step_name=next_step_name,
+            name=is_completed_listener_name,
+            dag=dag,
+        )
+        complete = build_complete(uid=uid or complete_name, dag=dag)
+        rejected = build_rejected(uid=uid or rejected_name, dag=dag)
+        decision = PythonOperator(
+            task_id=f"decision-{uid or name}",
+            python_callable=tri_branch_func,
+            op_kwargs={
+                "pull_from": is_completed_listener_name,
+                "negative_task": rejected_name,
+                "neutral_task": next_step,  # neutral_step_name in DAG
+                "positive_task": complete_name,
+            },
+            dag=dag,
+        )
+
+        is_completed >> decision >> [complete, next_step, rejected]
+
+        return task_group