Quickfix/bumping tenacity version

.github/workflows/ci-data.yml

@@ -17,10 +17,10 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Set up Python 3.7
+      - name: Set up Python 3.9
         uses: actions/setup-python@v2
         with:
-          python-version: 3.7
+          python-version: 3.9
 
       - name: Install dependencies
         run: |

Makefile

@@ -96,15 +96,15 @@ install: clean ## install the package to the active Python's site-packages
 
 
 install-dev: clean ## install the package to the active Python's site-packages
-	virtualenv -p python3 venv; \
+	virtualenv -p python3.9 venv; \
 	source venv/bin/activate; \
 	python -m pip install --upgrade pip; \
 	python setup.py install; \
 	pip install -e . ; \
-	pip install -r reqs/dev.txt -r reqs/test.txt
+	SYSTEM_VERSION_COMPAT=0 CFLAGS='-std=c++20' pip install -r reqs/dev.txt -r reqs/test.txt
 
 install-test: clean ## install the package to the active Python's site-packages
-	virtualenv -p python3 venv; \
+	virtualenv -p python3.9 venv; \
 	source venv/bin/activate; \
 	python -m pip install --upgrade pip; \
 	pip install -r reqs/test.txt -r reqs/base.txt

README.md

@@ -87,6 +87,41 @@ flowchart TD;
 
 ```
 
+Plugins for dagger
+-------
+
+### Overview
+Dagger now supports a plugin system that allows users to extend its functionality by adding custom Python classes. These plugins are integrated into the Jinja2 templating engine, enabling dynamic rendering of task configuration templates.
+### Purpose
+The plugin system allows users to define Python classes that can be loaded into the Jinja2 environment. When functions from these classes are invoked within a task configuration template, they are rendered dynamically using Jinja2. This feature enhances the flexibility of task configurations by allowing custom logic to be embedded directly in the templates.
+
+### Usage
+1. **Creating a Plugin:** To create a new plugin, define a Python class in a folder(for example `plugins/sample_plugin/sample_plugin.py`) with the desired methods. For example:
+```python
+class MyCustomPlugin:
+    def generate_input(self, branch_name):
+        return [{"name": f"{branch_name}", "type": "dummy"}]
+```
+This class defines a `generate_input` method that takes the branch_name from the module config and returns a dummy dagger task.
+
+2. **Loading the Plugin into Dagger:** To load this plugin into Dagger's Jinja2 environment, you need to register it in your `dagger_config.yaml`:
+```yaml
+# pipeline.yaml
+plugin:
+  paths: 
+    - plugins # all Python classes within this path will be loaded into the Jinja environment
+```
+
+3. **Using Plugin Methods in Templates:** Once the plugin is loaded, you can call its methods from within any Jinja2 template in your task configurations:
+```yaml
+# task_configuration.yaml
+type: batch
+description: sample task
+inputs:                        # format: list | Use dagger init-io cli
+  {{ MyCustomPlugin.generate_input("dummy_input") }}
+```
+
+
 
 Credits
 -------

dagger/cli/module.py

@@ -1,17 +1,34 @@
 import click
 from dagger.utilities.module import Module
 from dagger.utils import Printer
+import json
 
 
+def parse_key_value(ctx, param, value):
+    #print('YYY', value)
+    if not value:
+        return {}
+    key_value_dict = {}
+    for pair in value:
+        try:
+            key, val_file_path = pair.split('=', 1)
+            #print('YYY', key, val_file_path, pair)
+            val = json.load(open(val_file_path))
+            key_value_dict[key] = val
+        except ValueError:
+            raise click.BadParameter(f"Key-value pair '{pair}' is not in the format key=value")
+    return key_value_dict
+
 @click.command()
 @click.option("--config_file", "-c", help="Path to module config file")
 @click.option("--target_dir", "-t", help="Path to directory to generate the task configs to")
-def generate_tasks(config_file: str, target_dir: str) -> None:
+@click.option("--jinja_parameters", "-j", callback=parse_key_value, multiple=True, default=None, help="Path to jinja parameters json file in the format: <jinja_variable_name>=<path to json file>")
+def generate_tasks(config_file: str, target_dir: str, jinja_parameters: dict) -> None:
     """
     Generating tasks for a module based on config
     """
 
-    module = Module(config_file, target_dir)
+    module = Module(config_file, target_dir, jinja_parameters)
     module.generate_task_configs()
 
     Printer.print_success("Tasks are successfully generated")

dagger/conf.py

@@ -21,9 +21,7 @@
 # Airflow parameters
 airflow_config = config.get('airflow', None) or {}
 WITH_DATA_NODES = airflow_config.get('with_data_nodes', False)
-EXTERNAL_SENSOR_POKE_INTERVAL = airflow_config.get('external_sensor_poke_interval', 600)
-EXTERNAL_SENSOR_TIMEOUT = airflow_config.get('external_sensor_timeout', 28800)
-EXTERNAL_SENSOR_MODE = airflow_config.get('external_sensor_mode', 'reschedule')
+EXTERNAL_SENSOR_DEFAULT_ARGS = airflow_config.get('external_sensor_default_args', {})
 IS_DUMMY_OPERATOR_SHORT_CIRCUIT = airflow_config.get('is_dummy_operator_short_circuit', False)
 
 # Neo4j parameters
@@ -100,4 +98,15 @@
 # Alert parameters
 alert_config = config.get('alert', None) or {}
 SLACK_TOKEN = alert_config.get('slack_token', None)
-DEFAULT_ALERT = alert_config.get('default_alert', {"type": "slack", "channel": "#airflow-jobs", "mentions": None})
+DEFAULT_ALERT = alert_config.get('default_alert', {"type": "slack", "channel": "#airflow-jobs", "mentions": None})
+
+# Plugin parameters
+plugin_config = config.get('plugin', None) or {}
+PLUGIN_DIRS = [os.path.join(AIRFLOW_HOME, path) for path in plugin_config.get('paths', [])]
+logging.info(f"All Python classes will be loaded as plugins from the following directories: {PLUGIN_DIRS}")
+
+# ReverseETL parameters
+reverse_etl_config = config.get('reverse_etl', None) or {}
+REVERSE_ETL_DEFAULT_JOB_NAME = reverse_etl_config.get('default_job_name', None)
+REVERSE_ETL_DEFAULT_EXECUTABLE_PREFIX = reverse_etl_config.get('default_executable_prefix', None)
+REVERSE_ETL_DEFAULT_EXECUTABLE = reverse_etl_config.get('default_executable', None)

dagger/dag_creator/airflow/dag_creator.py

@@ -58,7 +58,7 @@ def _get_external_task_sensor_name_dict(self, from_task_id: str) -> dict:
             "external_sensor_name": f"{from_pipeline_name}-{from_task_name}-sensor",
         }
 
-    def _get_external_task_sensor(self, from_task_id: str, to_task_id: str) -> ExternalTaskSensor:
+    def _get_external_task_sensor(self, from_task_id: str, to_task_id: str, follow_external_dependency: dict) -> ExternalTaskSensor:
         """
         create an object of external task sensor for a specific from_task_id and to_task_id
         """
@@ -72,6 +72,9 @@ def _get_external_task_sensor(self, from_task_id: str, to_task_id: str) -> Exter
 
         to_pipe_id = self._task_graph.get_node(to_task_id).obj.pipeline.name
 
+        extra_args = conf.EXTERNAL_SENSOR_DEFAULT_ARGS.copy()
+        extra_args.update(follow_external_dependency)
+
         return ExternalTaskSensor(
             dag=self._dags[to_pipe_id],
             task_id=external_sensor_name,
@@ -80,9 +83,7 @@ def _get_external_task_sensor(self, from_task_id: str, to_task_id: str) -> Exter
             execution_date_fn=self._get_execution_date_fn(
                 from_pipeline_schedule, to_pipeline_schedule
             ),
-            mode=conf.EXTERNAL_SENSOR_MODE,
-            poke_interval=conf.EXTERNAL_SENSOR_POKE_INTERVAL,
-            timeout=conf.EXTERNAL_SENSOR_TIMEOUT,
+            **extra_args
         )
 
     def _create_control_flow_task(self, pipe_id, dag):
@@ -135,6 +136,7 @@ def _create_edge_without_data(self, from_task_id: str, to_task_ids: list, node:
             to_task_ids: The IDs of the tasks to which the edge connects.
             node: The current node in a task graph.
         """
+
         from_pipe = (
             self._task_graph.get_node(from_task_id).obj.pipeline_name if from_task_id else None
         )
@@ -143,7 +145,7 @@ def _create_edge_without_data(self, from_task_id: str, to_task_ids: list, node:
             to_pipe = self._task_graph.get_node(to_task_id).obj.pipeline_name
             if from_pipe and from_pipe == to_pipe:
                 self._tasks[from_task_id] >> self._tasks[to_task_id]
-            elif from_pipe and from_pipe != to_pipe and edge_properties.follow_external_dependency:
+            elif from_pipe and from_pipe != to_pipe and edge_properties.follow_external_dependency is not None:
                 from_schedule = self._task_graph.get_node(from_task_id).obj.pipeline.schedule
                 to_schedule = self._task_graph.get_node(to_task_id).obj.pipeline.schedule
                 if not from_schedule.startswith("@") and not to_schedule.startswith("@"):
@@ -155,15 +157,17 @@ def _create_edge_without_data(self, from_task_id: str, to_task_ids: list, node:
                         not in self._sensor_dict.get(to_pipe, dict()).keys()
                     ):
                         external_task_sensor = self._get_external_task_sensor(
-                            from_task_id, to_task_id
+                            from_task_id, to_task_id, edge_properties.follow_external_dependency
                         )
-                        self._sensor_dict[to_pipe] = {
+
+                        if self._sensor_dict.get(to_pipe) is None:
+                            self._sensor_dict[to_pipe] = {}
+
+                        self._sensor_dict[to_pipe].update({
                             external_task_sensor_name: external_task_sensor
-                        }
-                        (
-                            self._tasks[self._get_control_flow_task_id(to_pipe)]
-                            >> external_task_sensor
-                        )
+                        })
+
+                        self._tasks[self._get_control_flow_task_id(to_pipe)] >> external_task_sensor
                     self._sensor_dict[to_pipe][external_task_sensor_name] >> self._tasks[to_task_id]
             else:
                 self._tasks[self._get_control_flow_task_id(to_pipe)] >> self._tasks[to_task_id]

dagger/dag_creator/airflow/operator_creator.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from datetime import timedelta
+from airflow.utils.task_group import TaskGroup
 
 TIMEDELTA_PARAMETERS = ['execution_timeout']
 
@@ -11,6 +12,15 @@ def __init__(self, task, dag):
         self._template_parameters = {}
         self._airflow_parameters = {}
 
+    def _get_existing_task_group_or_create_new(self):
+        group_id = self._task.task_group
+        if self._dag.task_group:
+            for group in self._dag.task_group.children.values():
+                if isinstance(group, TaskGroup) and group.group_id == group_id:
+                    return group
+
+        return TaskGroup(group_id=group_id, dag=self._dag)
+
     @abstractmethod
     def _create_operator(self, kwargs):
         raise NotImplementedError
@@ -34,6 +44,9 @@ def _update_airflow_parameters(self):
         if self._task.timeout_in_seconds:
             self._airflow_parameters["execution_timeout"] = self._task.timeout_in_seconds
 
+        if self._task.task_group:
+            self._airflow_parameters["task_group"] = self._get_existing_task_group_or_create_new()
+
         self._fix_timedelta_parameters()
 
     def create_operator(self):

dagger/dag_creator/airflow/operator_creators/batch_creator.py

@@ -1,5 +1,9 @@
+from pathlib import Path
+from datetime import timedelta
+
 from dagger.dag_creator.airflow.operator_creator import OperatorCreator
 from dagger.dag_creator.airflow.operators.awsbatch_operator import AWSBatchOperator
+from dagger import conf
 
 
 class BatchCreator(OperatorCreator):
@@ -8,6 +12,20 @@ class BatchCreator(OperatorCreator):
     def __init__(self, task, dag):
         super().__init__(task, dag)
 
+    @staticmethod
+    def _validate_job_name(job_name, absolute_job_name):
+        if not absolute_job_name and not job_name:
+            raise Exception("Both job_name and absolute_job_name cannot be null")
+
+        if absolute_job_name is not None:
+            return absolute_job_name
+
+        job_path = Path(conf.DAGS_DIR) / job_name.replace("-", "/")
+        assert (
+            job_path.is_dir()
+        ), f"Job name `{job_name}`, points to a non-existing folder `{job_path}`"
+        return job_name
+
     def _generate_command(self):
         command = [self._task.executable_prefix, self._task.executable]
         for param_name, param_value in self._template_parameters.items():
@@ -21,16 +39,16 @@ def _create_operator(self, **kwargs):
         overrides = self._task.overrides
         overrides.update({"command": self._generate_command()})
 
+        job_name = self._validate_job_name(self._task.job_name, self._task.absolute_job_name)
         batch_op = AWSBatchOperator(
             dag=self._dag,
             task_id=self._task.name,
-            job_name=self._task.job_name,
-            absolute_job_name=self._task.absolute_job_name,
+            job_name=self._task.name,
+            job_definition=job_name,
             region_name=self._task.region_name,
-            cluster_name=self._task.cluster_name,
             job_queue=self._task.job_queue,
-            overrides=overrides,
+            container_overrides=overrides,
+            awslogs_enabled=True,
             **kwargs,
         )
-
         return batch_op

dagger/dag_creator/airflow/operator_creators/dbt_creator.py

@@ -13,28 +13,28 @@ def __init__(self, task, dag):
         self._project_dir = task.project_dir
         self._profile_dir = task.profile_dir
         self._profile_name = task.profile_name
+        self._target_name = task.target_name
         self._select = task.select
         self._dbt_command = task.dbt_command
+        self._vars = task.vars
+        self._create_external_athena_table = task.create_external_athena_table
 
     def _generate_command(self):
         command = [self._task.executable_prefix, self._task.executable]
         command.append(f"--project_dir={self._project_dir}")
         command.append(f"--profiles_dir={self._profile_dir}")
         command.append(f"--profile_name={self._profile_name}")
+        command.append(f"--target_name={self._target_name}")
         command.append(f"--dbt_command={self._dbt_command}")
         if self._select:
             command.append(f"--select={self._select}")
-
-        if len(self._template_parameters) > 0:
-            dbt_vars = json.dumps(self._template_parameters)
+        if self._vars:
+            dbt_vars = json.dumps(self._vars)
             command.append(f"--vars='{dbt_vars}'")
-
+        if self._create_external_athena_table:
+            command.append(f"--create_external_athena_table={self._create_external_athena_table}")
+        for param_name, param_value in self._template_parameters.items():
+            command.append(
+                f"--{param_name}={param_value}"
+            )
         return command
-
-    # Overwriting function because for dbt we don't want to add inputs/outputs to the
-    # template parameters.
-    def create_operator(self):
-        self._template_parameters.update(self._task.template_parameters)
-        self._update_airflow_parameters()
-
-        return self._create_operator(**self._airflow_parameters)

dagger/dag_creator/airflow/operator_creators/redshift_load_creator.py

@@ -2,9 +2,8 @@
 from typing import Optional
 
 from dagger.dag_creator.airflow.operator_creator import OperatorCreator
-from dagger.dag_creator.airflow.operators.redshift_sql_operator import (
-    RedshiftSQLOperator,
-)
+from dagger.dag_creator.airflow.operators.redshift_sql_operator import RedshiftSQLOperator
+
 
 
 class RedshiftLoadCreator(OperatorCreator):

dagger/dag_creator/airflow/operator_creators/redshift_transform_creator.py

@@ -1,7 +1,7 @@
 from os.path import join
 
 from dagger.dag_creator.airflow.operator_creator import OperatorCreator
-from dagger.dag_creator.airflow.operators.postgres_operator import PostgresOperator
+from dagger.dag_creator.airflow.operators.redshift_sql_operator import RedshiftSQLOperator
 
 
 class RedshiftTransformCreator(OperatorCreator):
@@ -22,11 +22,11 @@ def _read_sql(directory, file_path):
     def _create_operator(self, **kwargs):
         sql_string = self._read_sql(self._task.pipeline.directory, self._task.sql_file)
 
-        redshift_op = PostgresOperator(
+        redshift_op = RedshiftSQLOperator(
             dag=self._dag,
             task_id=self._task.name,
             sql=sql_string,
-            postgres_conn_id=self._task.postgres_conn_id,
+            redshift_conn_id=self._task.postgres_conn_id,
             params=self._template_parameters,
             **kwargs,
         )

dagger/dag_creator/airflow/operator_creators/redshift_unload_creator.py

@@ -1,7 +1,7 @@
 from os.path import join
 
 from dagger.dag_creator.airflow.operator_creator import OperatorCreator
-from dagger.dag_creator.airflow.operators.postgres_operator import PostgresOperator
+from dagger.dag_creator.airflow.operators.redshift_sql_operator import RedshiftSQLOperator
 
 REDSHIFT_UNLOAD_CMD = """
 unload ('{sql_string}')
@@ -58,12 +58,13 @@ def _create_operator(self, **kwargs):
 
         unload_cmd = self._get_unload_command(sql_string)
 
-        redshift_op = PostgresOperator(
+        redshift_op = RedshiftSQLOperator(
             dag=self._dag,
             task_id=self._task.name,
             sql=unload_cmd,
-            postgres_conn_id=self._task.postgres_conn_id,
+            redshift_conn_id=self._task.postgres_conn_id,
             params=self._template_parameters,
+            autocommit=True,
             **kwargs,
         )