regscale-cli 6.16.0.0__py3-none-any.whl

regscale/airflow/hierarchy.py

@@ -0,0 +1,88 @@
+"""Generate a hierarchy for Airflow based on Click."""
+
+from typing import Optional
+
+from regscale.regscale import cli
+from regscale.models.click_models import ClickCommand
+from regscale.airflow.click_mixins import AirflowClickGroup
+
+AIRFLOW_CLICK_GROUP = AirflowClickGroup.from_group(cli, prefix="regscale")
+AIRFLOW_CLICK_OPERATORS = AIRFLOW_CLICK_GROUP.flatten_operator()
+
+
+def iter_operators(operators: dict = AIRFLOW_CLICK_OPERATORS) -> dict:
+    """Iterate over the operators and create a structured dict"""
+
+    def _nested_dict(data: tuple, value: ClickCommand, result: dict = None) -> Optional[dict]:
+        """Create a single dict from a nested dict."""
+        if result is None:
+            result = {}
+        if len(data) > 1:
+            key = data[0]
+            if key not in result:
+                result[key] = {}
+            _nested_dict(data=data[1:], value=value, result=result[key])
+        else:
+            if data[0] in result:
+                if isinstance(result[data[0]], dict):
+                    result[data[0]].update(_make_platform_params(value))
+                else:
+                    result[data[0]] = [result[data[0]], _make_platform_params(value)]
+            else:
+                result[data[0]] = _make_platform_params(value)
+        return result
+
+    def _make_platform_params(command: ClickCommand) -> dict:
+        """Create a dict from a command of what platform would need."""
+        return {
+            "name": command.name,
+            "params": command.parameters,
+            "parameters": {**command.params},
+        }
+
+    new_dict = {}
+    for operator, op_dict in operators.items():
+        if "__" in operator:
+            operator_split = operator.split("__")
+            new_dict = _nested_dict(data=operator_split, value=op_dict["command"], result=new_dict)
+        else:
+            if operator in new_dict:
+                if isinstance(new_dict[operator], dict):
+                    new_dict[operator].update(_make_platform_params(op_dict["command"]))
+                else:
+                    new_dict[operator] = [
+                        new_dict[operator],
+                        _make_platform_params(op_dict["command"]),
+                    ]
+            else:
+                new_dict[operator] = _make_platform_params(op_dict["command"])
+    return new_dict
+
+
+AIRFLOW_CLI_HIERARCHY = iter_operators()
+
+
+def get_command_params(cmd: str) -> list:
+    """Get the parameters for a command
+
+    :param str cmd: the command to get the parameters for
+    :return: a list of parameters
+    :rtype: list
+    """
+    cmd_split = cmd.split(" ")
+    if len(cmd_split) == 1:
+        return AIRFLOW_CLI_HIERARCHY[cmd_split[0]]["params"]
+    elif len(cmd_split) == 2:
+        return AIRFLOW_CLI_HIERARCHY[cmd_split[0]][cmd_split[1]]["params"]
+
+
+def return_all_commands_and_params() -> dict:
+    """Return all commands and parameters
+
+    :return: a dict of commands and parameters
+    :rtype: dict
+    """
+    dct = {}
+    for operator in AIRFLOW_CLICK_OPERATORS:
+        click_name = " ".join(operator.split("__"))
+        dct[click_name] = {"parameters": AIRFLOW_CLICK_OPERATORS[operator]["command"].parameters}

regscale/airflow/operators/click.py

@@ -0,0 +1,36 @@
+"""Generate PythonOperators based on a click.Command."""
+
+from typing import Union, Optional
+
+import click
+from airflow.operators.python import PythonOperator
+
+from airflow import DAG
+from regscale.airflow.tasks.click import execute_click_command
+from regscale.models.click_models import ClickCommand
+
+
+def generate_operator_for_command(
+    command: Union[click.Command, ClickCommand],
+    command_name: Optional[str] = None,
+    dag: Optional[DAG] = None,
+    **kwargs: dict,
+) -> PythonOperator:
+    """Generate a PythonOperator for a Click Command
+
+    :param Union[click.Command, ClickCommand] command: the command to generate the operator for
+    :param Optional[str] command_name: optional command name to specify as task_id, defaults to None
+    :param Optional[DAG] dag: an Optional airflow DAG to pass, defaults to None
+    :param dict **kwargs: additional named parameters to pass to PythonOperator
+    :return: a PythonOperator configured to execute the Click Command
+    :rtype: PythonOperator
+    """
+    if isinstance(command, click.Command):
+        command = ClickCommand.from_command(command)
+    return PythonOperator(
+        task_id=command_name or command.name,
+        python_callable=execute_click_command,
+        op_kwargs={"command": command.callback},
+        dag=dag,
+        **kwargs,
+    )

regscale/airflow/sensors/sql.py

@@ -0,0 +1,107 @@
+"""SQL Sensors for Airflow"""
+
+from typing import Literal
+from uuid import uuid4
+
+from airflow.sensors.sql import SqlSensor
+
+from regscale.airflow.sessions.sql import SQLQuery
+
+
+class SqlSensorXcon(SqlSensor):
+    def __init__(self, step_name: str = "sql_sensor_xcon", **kwargs):
+        super().__init__(**kwargs)
+        self.step_name = step_name
+
+    def poke(self, context):
+        self.log.info(f"Poking: {self.sql}")
+        records = self._get_hook().get_first(self.sql, parameters=self.parameters)
+        if not records:
+            return False
+        context["task_instance"].xcom_push(key=self.step_name, value=records)
+        return True
+
+
+def build_sql_sensor_xcon(
+    sql: SQLQuery,
+    name: str = None,
+    conn_id: str = "platform_db",
+    unique: bool = False,
+    mode: Literal["reschedule", "poke"] = "poke",
+    step_name: str = "sql_sensor_xcon",
+    timeout: int = 60 * 60 * 24 * 7,
+    poke_interval: int = 60 * 5,
+    **kwargs: dict,
+) -> SqlSensorXcon:
+    """Build a SQL Sensor that will push its results to XCom
+
+    :param SQLQuery sql: the SQL query to run
+    :param str name: the name of the sensor, defaults to None
+    :param str conn_id: the connection id to use, defaults to "platform_db"
+    :param bool unique: whether to make the sensor unique, defaults to False
+    :param Literal["reschedule", "poke"] mode: the mode of the sensor, defaults to "poke"
+    :param str step_name: the name of the step, defaults to "sql_sensor_xcon"
+    :param int timeout: the timeout of the sensor, defaults to 60 * 60 * 24 * 7
+    :param int poke_interval: the poke interval of the sensor, defaults to 60 * 5
+    :param dict kwargs: keyword arguments to pass to the sensor
+    :return: the SQL Sensor
+    :rtype: SqlSensorXcon
+    """
+    dag_run_dict = dict(zip(sql.dag_run_vars, sql.dag_run_vars))
+    params_dict = dict(zip(sql.params, sql.params))
+    if unique and name is not None:
+        name = f"{name}-{str(uuid4())[:8]}"
+
+    return SqlSensorXcon(
+        task_id=name or f"sql_sensor-{name or str(uuid4())[:8]}",
+        sql=sql.raw,
+        conn_id=conn_id,
+        params=params_dict or dag_run_dict,
+        mode=mode,
+        step_name=step_name,
+        timeout=timeout,
+        poke_interval=poke_interval,
+        **kwargs,
+    )
+
+
+def build_sql_sensor(
+    sql: SQLQuery,
+    name: str = None,
+    conn_id: str = "platform_db",
+    unique: bool = False,
+    mode: Literal["reschedule", "poke"] = "poke",
+    timeout: int = 60 * 60 * 24 * 7,
+    poke_interval: int = 60 * 5,
+    **kwargs: dict,
+) -> SqlSensor:
+    """Build a SQL Sensor from a SQLQuery object
+
+    :param SQLQuery sql: the SQL query to run
+    :param str name: the name of the sensor, defaults to None
+    :param str conn_id: the connection id to use, defaults to "platform_db"
+    :param bool unique: whether to make the sensor unique, defaults to False
+    :param Literal["reschedule", "poke"] mode: the mode of the sensor, defaults to "poke"
+    :param int timeout: the timeout of the sensor, defaults to 60 * 60 * 24 * 7
+    :param int poke_interval: the poke interval of the sensor, defaults to 60 * 5
+    :param dict **kwargs: additional keyword arguments
+    :return: the SQL Sensor
+    :rtype: SqlSensor
+    """
+    dag_run_dict = dict(zip(sql.dag_run_vars, sql.dag_run_vars))
+    params_dict = dict(zip(sql.params, sql.params))
+    if unique and name:
+        name = f"{name}-{str(uuid4())[:8]}"
+    if unique and name is None:
+        name = f"sql_sensor-{str(uuid4())[:8]}"
+
+    return SqlSensor(
+        task_id=name,
+        sql=sql.raw,
+        conn_id=conn_id,
+        params=params_dict or dag_run_dict,
+        mode=mode,
+        timeout=timeout,
+        poke_interval=poke_interval,
+        **kwargs,
+    )

regscale/airflow/sessions/sql/__init__.py

@@ -0,0 +1,3 @@
+"""Import"""
+
+from regscale.airflow.sessions.sql.queries import SQLQuery

regscale/airflow/sessions/sql/queries.py

@@ -0,0 +1,64 @@
+"""Provide SQL Queries to be used in Airflow DAGs related to SQL Server."""
+
+from typing import List
+
+from regscale.utils.string import (
+    replace_unknown_spaces_and_newlines,
+    extract_double_curly_braces_contents,
+    extract_dag_run_conf_key,
+    extract_param,
+)
+
+
+class SQLQuery:
+    """
+    A class to represent a SQL Query
+
+    :param str query: a SQL query
+    :param dict **kwargs: additional keyword arguments
+    """
+
+    def __init__(self, query: str = "", **kwargs: dict):
+        self.raw = query
+        self.query = replace_unknown_spaces_and_newlines(query)
+        self.kwargs = kwargs
+        self.vars: List[str] = []
+        self.dag_run_vars: List[str] = []
+        self.params: List[str] = []
+        self.jinja_kwargs = {}
+        self._populate_jinja_vars(**kwargs)
+
+    def __repr__(self) -> str:
+        """Return the SQLQuery object
+
+        :return: the SQLQuery object as a string
+        :rtype: str
+        """
+        return f"SQLQuery(query={self.query}, kwargs={self.kwargs})"
+
+    def __str__(self) -> str:
+        """Return the SQLQuery object
+
+        :return: the SQLQuery object as a string
+        :rtype: str
+        """
+        return f"{self.dag_run_vars=}\n{self.params=}\n{self.unique_vars=}"
+
+    def _populate_jinja_vars(self, **kwargs: dict):
+        """Populate the SQLQuery object with Jinja variables
+
+        :param dict **kwargs: keyword arguments
+        """
+        self.jinja_kwargs = kwargs
+        self.vars = extract_double_curly_braces_contents(self.query)
+        self.dag_run_vars = [extract_dag_run_conf_key(__) for __ in self.vars]
+        self.params = [extract_param(_) for _ in self.vars]
+
+    @property
+    def unique_vars(self) -> List[str]:
+        """Return a list of unique Jinja variables
+
+        :return: a list of unique Jinja variables as strings
+        :rtype: List[str]
+        """
+        return list(set(self.dag_run_vars + self.params))

regscale/airflow/sessions/sql/sql_server_queries.py

@@ -0,0 +1,248 @@
+"""SQL Server Platform Queries."""
+
+from regscale.airflow.sessions.sql.queries import SQLQuery
+
+
+CHECK_WORKFLOW_STATUS_QUERY: str = """
+SELECT
+  CASE
+    WHEN EXISTS (
+      SELECT 1
+      FROM dbo.WorkflowInstance wi
+      WHERE wi.id = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+      AND NOT EXISTS (
+        SELECT 1
+        FROM dbo.WorkflowInstanceSteps wis
+        WHERE wis.WorkflowInstanceId = wi.id
+        AND wis.[Status] <> 'Rejected'
+      )
+    ) THEN -1
+    WHEN EXISTS (
+      SELECT 1
+      FROM dbo.WorkflowInstance wi
+      WHERE wi.id = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+      AND NOT EXISTS (
+        SELECT 1
+        FROM dbo.WorkflowInstanceSteps wis
+        WHERE wis.WorkflowInstanceId = wi.id
+        AND wis.[Status] <> 'Approved'
+      )
+    ) THEN 1
+    WHEN EXISTS(
+      SELECT 1
+      FROM dbo.WorkflowInstance wi
+      WHERE wi.id = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+      AND NOT EXISTS (
+        SELECT 1
+        FROM dbo.WorkflowInstanceSteps wis
+        WHERE wis.WorkflowInstanceId = wi.id
+        AND wis.[Status] <> 'Pending'
+    ) THEN 0
+    ELSE NULL
+  END AS OverallStatus;
+"""
+
+NEW_CHECK_WORKFLOW_STATUS_QUERY: str = """
+declare @workflowInstanceId int = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+declare @workflowStepId int = -1
+;WITH orderedWorkFlowSteps AS
+(
+   SELECT
+       id, ROW_NUMBER() OVER(ORDER BY [Order] DESC) AS 'RowNum'
+   FROM WorkflowInstanceSteps where WorkflowInstanceId = @workflowInstanceId
+)
+SELECT @workflowStepId = id
+FROM orderedWorkFlowSteps
+WHERE RowNum = 2
+SELECT
+    case
+        when exists
+            (select wis.id from dbo.WorkflowInstanceSteps wis
+             where wis.Id = @workflowStepId and wis.WorkflowInstanceId = @workflowInstanceId and wis.[Status] = 'Rejected')
+            then -1
+        when exists
+            (select wis.id from dbo.WorkflowInstanceSteps wis
+             where wis.Id = @workflowStepId and wis.WorkflowInstanceId = @workflowInstanceId and wis.[Status] = 'Approved')
+            then 1
+        when exists
+            (select wis.id from dbo.WorkflowInstanceSteps wis
+             where wis.Id = @workflowStepId and wis.WorkflowInstanceId = @workflowInstanceId and wis.[Status] = 'Pending')
+            then 0
+        else
+            null
+        end as  OverallStatus;
+"""
+
+NEWER_CHECK_WORKFLOW_STATUS_QUERY: str = """
+declare @workflowInstanceId int = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+;select * from WorkflowInstanceSteps
+where WorkflowInstanceId = @workflowInstanceId and [Order] not in (0, 1000) order by [order] -- this is for validating results only
+select top 1
+case when wis.status = 'Rejected' then -1
+        when wis.status = 'Approved' then 1
+        when wis.status = 'Pending' then 0
+    else null
+    end as OverallStatus,
+ *
+from WorkflowInstanceSteps wis
+where wis.WorkflowInstanceId = @workflowInstanceId
+and status in ('Rejected', 'Approved', 'Pending' )
+order by [Order] DESC;
+"""
+
+CHECK_IF_COMPLETED_QUERY: str = """
+DECLARE @workflowInstanceId INT = '{{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}';
+WITH OrderedSteps AS (
+    SELECT
+        wis.*,
+        LEAD(wis.[Order]) OVER (ORDER BY wis.[Order] DESC) AS NextOrder,
+        LEAD(wis.[Status]) OVER (ORDER BY wis.[Order] DESC) AS NextStatus
+    FROM WorkflowInstanceSteps wis
+    WHERE wis.WorkflowInstanceId = @workflowInstanceId
+)
+SELECT
+    CASE
+        WHEN [Order] = 1000 AND NextStatus = 'Approved' THEN 1
+        ELSE NULL
+    END AS Result
+FROM OrderedSteps
+WHERE [Order] = 1000;
+"""
+
+WATCH_WORFKLOW_QUERY: str = """
+DECLARE @workflowInstanceId INT = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }};
+WITH OrderedSteps AS (
+    SELECT
+        wis.*,
+        LEAD(wis.[Order]) OVER (ORDER BY wis.[Order] DESC) AS NextOrder,
+        ROW_NUMBER() OVER (ORDER BY wis.[Order] DESC) AS Rn
+    FROM WorkflowInstanceSteps wis
+    WHERE wis.WorkflowInstanceId = @workflowInstanceId
+        AND status IN ('Rejected', 'Approved', 'Pending')
+)
+SELECT TOP 1
+    CASE
+        WHEN NextOrder = 1000 THEN 1
+        ELSE 0
+    END AS Result,
+    wis.*
+FROM OrderedSteps wis
+WHERE Rn = 1;
+"""
+
+GET_WORKFLOW_STEPS_QUERY: str = """
+SELECT COUNT(*) as Steps
+FROM WorkflowInstanceSteps as wis
+WHERE wis.WorkflowInstanceId = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+GROUP BY WorkflowInstanceId;
+"""
+
+BETTER_GET_WORKFLOW_STEPS_QUERY: str = """
+declare @workflowInstanceId int = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }}
+;select
+    stepCount.Steps,
+FROM
+    (
+    SELECT
+        @workflowInstanceId WorkflowinstanceId) as t1
+left JOIN (
+    SELECT
+        @workflowInstanceId ID,
+        COUNT(*) as Steps
+    FROM
+        WorkflowInstanceSteps as wis
+    WHERE
+        wis.WorkflowInstanceId = @workflowInstanceId
+    group by
+        WorkflowInstanceId) as stepCount on
+    t1.workflowInstanceId = stepCount.ID
+LEFT JOIN (
+    select
+        top 1
+wis.WorkflowinstanceId,
+        case
+            when wis.status = 'Rejected' then -1
+            when wis.status = 'Approved' then 1
+            when wis.status = 'Pending' then 0
+            else null
+        end as OverallStatus
+    from
+        WorkflowInstanceSteps as wis
+    where
+        wis.WorkflowInstanceId = @workflowInstanceId
+        and status in ('Rejected', 'Approved', 'Pending' )
+    order by
+        [Order] DESC) as status on
+    t1.WorkflowinstanceId = status.WorkflowinstanceId
+"""
+
+GET_STATUS_WORKFLOW_QUERY: str = """
+declare @workflowInstanceId int = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }};
+select
+    t1.WorkflowinstanceId,
+    stepCount.Steps,
+    status.OverallStatus
+FROM
+    (
+    SELECT
+        @workflowInstanceId WorkflowinstanceId) as t1
+left JOIN (
+    SELECT
+        @workflowInstanceId ID,
+        COUNT(*) as Steps
+    FROM
+        WorkflowInstanceSteps as wis
+    WHERE
+        wis.WorkflowInstanceId = @workflowInstanceId
+    group by
+        WorkflowInstanceId) as stepCount on
+    t1.workflowInstanceId = stepCount.ID
+LEFT JOIN (
+    select
+        top 1
+wis.WorkflowinstanceId,
+        case
+            when wis.status = 'Rejected' then -1
+            when wis.status = 'Approved' then 1
+            when wis.status = 'Pending' then 0
+            else null
+        end as OverallStatus
+    from
+        WorkflowInstanceSteps as wis
+    where
+        wis.WorkflowInstanceId = @workflowInstanceId
+        and status in ('Rejected', 'Approved', 'Pending' )
+    order by
+        [Order] DESC) as status on
+    t1.WorkflowinstanceId = status.WorkflowinstanceId;
+"""
+
+GET_STEPS_QUERY: str = """
+declare @workflowInstanceId int = {{ dag_run.conf['workflow_id'] if dag_run and dag_run.conf and 'workflow_id' in dag_run.conf else params.workflow_id }};
+select
+    stepCount.Steps
+FROM
+    (
+    SELECT
+        @workflowInstanceId WorkflowinstanceId) as t1
+left JOIN (
+    SELECT
+        @workflowInstanceId ID,
+        COUNT(*) as Steps
+    FROM
+        WorkflowInstanceSteps as wis
+    WHERE
+        wis.WorkflowInstanceId = @workflowInstanceId
+    group by
+        WorkflowInstanceId) as stepCount on
+    t1.workflowInstanceId = stepCount.ID;
+"""
+
+NEW_WORKFLOW_SQL_QUERY: SQLQuery = SQLQuery(query=NEWER_CHECK_WORKFLOW_STATUS_QUERY)
+WORKFLOW_SQL_QUERY: SQLQuery = SQLQuery(query=NEW_CHECK_WORKFLOW_STATUS_QUERY)
+CHECK_IF_COMPLETED_SQL_QUERY: SQLQuery = SQLQuery(query=CHECK_IF_COMPLETED_QUERY)
+WATCH_WORKFLOW_SQL_QUERY: SQLQuery = SQLQuery(query=WATCH_WORFKLOW_QUERY)
+GET_WORKFLOW_STEPS_SQL_QUERY: SQLQuery = SQLQuery(query=GET_WORKFLOW_STEPS_QUERY)
+BETTER_GET_WORFKLOW_STEPS_SQL_QUERY: SQLQuery = SQLQuery(query=BETTER_GET_WORKFLOW_STEPS_QUERY)
+GET_STATUS_WORKFLOW_SQL_QUERY: SQLQuery = SQLQuery(query=GET_STATUS_WORKFLOW_QUERY)
+GET_STEPS_SQL_QUERY: SQLQuery = SQLQuery(query=GET_STEPS_QUERY)

regscale/airflow/tasks/branches.py

@@ -0,0 +1,22 @@
+"""Provide information related to branches in python."""
+
+
+def tri_branch_func(pull_from: str, negative_task: str, neutral_task: str, positive_task: str, **kwargs: dict) -> str:
+    """Return the branch name
+
+    :param str pull_from: the task_id to pull from
+    :param str negative_task: the task_id to go to if the pull_from task returns -1
+    :param str neutral_task: the task_id to go to if the pull_from task returns 0
+    :param str positive_task: the task_id to go to if the pull_from task returns 1
+    :param dict **kwargs: additional keyword arguments
+    :return: the task_id to go to
+    :rtype: str
+    """
+    ti = kwargs["ti"]
+    var = ti.xcom_pull(task_ids=pull_from)
+    if var == -1:
+        return negative_task
+    elif var == 0:
+        return neutral_task
+    else:
+        return positive_task