Feature/lightbeam delete dag

README.md

@@ -509,6 +509,47 @@ If you are using XComs in this manner, it is recommended to use the optional `gr
 your task-group to ensure that its value is static and easily-referenced in your DAG-initialization code.
 
 
+
+-----
+
+# LightbeamDeleteDAG
+`LightbeamDeleteDAG` is an Airflow DAG that can be used to delete data from an Ed-FI ODS using the `LightbeamOperator`.
+
+To avoid accidental deletions of entire resources, the `query_paramters` DAG-level config is required. 
+There is also an optional DAG-level config for the list of endpoints to target for deletion (defaults to `assessments`, `objectiveAssessments`, and `studentAssessments`).
+
+<details>
+<summary>DAG-level Arguments:</summary>
+
+| Argument        | Description                                                                                                    |
+|:----------------|:---------------------------------------------------------------------------------------------------------------|
+| lightbeam_path  | Path to installed Lightbeam package                                                                            |
+| pool            | Airflow pool against which all operations are applied                                                          |
+| fast_cleanup    | Boolean flag for whether to remove local files immediately upon failure, or only after success (default False) |
+
+Additional `EACustomDAG` parameters (e.g. `slack_conn_id`, `schedule_interval`, `default_args`, etc.) can be passed as kwargs.
+
+-----
+
+</details>
+
+`LightbeamDeleteDAG` builds task-groups for each (tenant, year). Task-groups are fully independent of one another.
+
+<details>
+<summary>Taskgroup-level Arguments:</summary>
+
+| Argument                     | Description                                                                                                                    |
+|:-----------------------------|:-------------------------------------------------------------------------------------------------------------------------------|
+| tenant_code                  | ODS-tenant representation to be saved in Snowflake tables                                                                      |
+| api_year                     | ODS API-year to be saved in Snowflake tables                                                                                   |
+| edfi_conn_id                 | Airflow connection with Ed-Fi ODS credentials and metadata defined for a specific tenant                                       |
+| lightbeam_kwargs             | Kwarg command-line arguments to be passed into `LightbeamOperator`                                                             |
+
+-----
+
+</details>
+
+
 -----
 
 ## Operators
@@ -548,7 +589,7 @@ Extends `BashOperator` to run Lightbeam with optional CLI arguments.
 | Argument            | Description                                                                                                                       |
 |:--------------------|:----------------------------------------------------------------------------------------------------------------------------------|
 | lightbeam_path      | Path to installed Ligthbeam package                                                                                               | 
-| command             | Lightbeam run command (i.e., `send`, `send+validate`, `validate`, `delete`) (default `send`)                                      | 
+| command             | Lightbeam run command (i.e., `send`, `send+validate`, `validate`, `fetch`, `delete`) (default `send`)                             | 
 | data_dir            | Directory of files for Lightbeam to use (also definable within the config file)                                                   |  
 | state_dir           | Path to directory where Lightbeam saves state between runs  (also definable within the config file)                               |
 | edfi_conn_id        | Optional Airflow connection with Ed-Fi ODS credentials (if present, fills config environment variables prefixed with `EDFI_API_`) |
@@ -560,6 +601,9 @@ Extends `BashOperator` to run Lightbeam with optional CLI arguments.
 | force               | Boolean flag to force a Lightbeam run, even if payloads have already been sent (default `False`)                                  |
 | older_than          | Optional timestamp string to filter payloads against when re-running Lightbeam                                                    |
 | newer_than          | Optional timestamp string to filter payloads against when re-running Lightbeam                                                    |
+| query               | Optional JSON payload to add query parameters to `fetch`commands                                                                  |
+| keep_keys           | Optional string to keep only specific keys from every `fetch`ed payload                                                           |
+| drop_keys           | Optional string to remove specific keys from every `fetch`ed payload                                                              |
 | resend_status_codes | Optional list of status-codes to filter payloads against when re-running Lightbeam                                                |
 
 -----

edu_edfi_airflow/__init__.py

@@ -1,5 +1,6 @@
 from edu_edfi_airflow.dags.earthbeam_dag import EarthbeamDAG
 from edu_edfi_airflow.dags.edfi_resource_dag import EdFiResourceDAG
+from edu_edfi_airflow.dags.lightbeam_delete_dag import LightbeamDeleteDAG
 from edu_edfi_airflow.providers.earthbeam.operators import EarthmoverOperator, LightbeamOperator
 from edu_edfi_airflow.providers.edfi.hooks.edfi import EdFiHook
 from edu_edfi_airflow.providers.edfi.transfers.edfi_to_s3 import EdFiToS3Operator

edu_edfi_airflow/dags/earthbeam_dag.py

@@ -1190,4 +1190,4 @@ def wrapper(*args, **kwargs):
                     )
 
             return result
-        return wrapper
+        return wrapper

edu_edfi_airflow/dags/lightbeam_delete_dag.py

@@ -0,0 +1,127 @@
+from typing import List, Optional
+
+from airflow.decorators import task, task_group
+from airflow.exceptions import AirflowFailException
+from airflow.models.param import Param
+
+import edfi_api_client
+from ea_airflow_util import EACustomDAG
+
+from edu_edfi_airflow.callables.s3 import remove_filepaths
+from edu_edfi_airflow.providers.earthbeam.operators import LightbeamOperator
+
+
+class LightbeamDeleteDAG:
+    """
+    Full Earthmover-Lightbeam DAG, with optional Python callable pre-processing.
+    """
+    emlb_state_directory  : str = '/efs/emlb/state'
+    emlb_results_directory: str = '/efs/emlb/results'
+    lb_output_directory   : str = '/efs/tmp_storage/lightbeam'
+
+    ENDPOINTS_LIST = [
+        "assessments",
+        "objectiveAssessments",
+        "studentAssessments",
+    ]
+
+    params_dict = {
+        "query_parameters": Param(
+            default={"exampleKey": "exampleValue"},
+            type="object",
+            description="JSON object of query parameters to define which records are to be deleted. At least one parameter is requied. For more information, see https://github.com/edanalytics/lightbeam?tab=readme-ov-file#fetch",
+        ),
+        "endpoints": Param(
+            default=ENDPOINTS_LIST,
+            type="array",
+            description="Newline-separated list of endpoints to delete"
+        ),
+    }
+
+    def __init__(self,
+        *,
+        lightbeam_path: Optional[str] = None,
+        pool: str = 'default_pool',
+        fast_cleanup: bool = False,
+
+        **kwargs
+    ):
+        self.lightbeam_path = lightbeam_path
+        self.pool = pool
+        self.fast_cleanup = fast_cleanup
+
+        self.dag = EACustomDAG(params=self.params_dict, **kwargs)
+
+
+    def build_tenant_year_taskgroup(self, 
+        tenant_code: str,
+        api_year: int,
+        edfi_conn_id: Optional[str] = None,
+        lightbeam_kwargs: Optional[dict] = None,
+        **kwargs
+    ):
+        """
+        Lightbeam: fetch -> Lightbeam: delete -> Clean-up
+
+        """
+
+        @task_group(prefix_group_id=True, group_id=f"{tenant_code}_{api_year}", dag=self.dag)
+        def tenant_year_taskgroup():
+
+            @task(multiple_outputs=True, pool=self.pool, dag=self.dag)
+            def run_lightbeam(command: str, **context):
+                if 'exampleKey' in context['params']['lightbeam_fetch_query']:
+                    raise AirflowFailException("No query parameters provided! At least one is required. This is a safety mechanism to prevent the deletion of entire resources.")
+
+                lightbeam_kwargs['query'] = context['params']['lightbeam_fetch_query']
+                lightbeam_kwargs['selector'] = context['params']['endpoints']
+
+                lb_output_dir = edfi_api_client.url_join(
+                    self.lb_output_directory,
+                    tenant_code, api_year,
+                    '{{ ds_nodash }}', '{{ ts_nodash }}'
+                )
+                lb_output_dir = context['task'].render_template(lb_output_dir, context)
+
+                lb_state_dir = edfi_api_client.url_join(
+                    self.emlb_state_directory,
+                    tenant_code, api_year,
+                    'lightbeam'
+                )
+
+                run_lightbeam = LightbeamOperator(
+                    task_id=f"run_lightbeam",
+                    lightbeam_path=self.lightbeam_path,
+                    data_dir=lb_output_dir,
+                    state_dir=lb_state_dir,
+                    edfi_conn_id=edfi_conn_id,
+                    **(lightbeam_kwargs or {}),
+                    command=command,
+                    dag=self.dag
+                )
+
+                return {
+                    "data_dir": run_lightbeam.execute(**context),
+                    "state_dir": lb_state_dir
+                }
+
+            @task(trigger_rule="all_done" if self.fast_cleanup else "all_success", pool=self.pool, dag=self.dag)
+            def remove_files(filepaths):
+                unnested_filepaths = []
+                for filepath in filepaths:
+                    if isinstance(filepath, str):
+                        unnested_filepaths.append(filepath)
+                    else:
+                        unnested_filepaths.extend(filepath)
+
+                return remove_filepaths(unnested_filepaths)
+
+            lightbeam_fetch = run_lightbeam.override(task_id="run_lightbeam_fetch")(command="fetch")
+            lightbeam_delete = run_lightbeam.override(task_id="run_lightbeam_delete")(command="delete")
+
+            # Final cleanup (apply at very end of the taskgroup)
+            remove_files_operator = remove_files(lightbeam_fetch["data_dir"])
+
+            lightbeam_fetch >> lightbeam_delete >> remove_files_operator
+
+        return tenant_year_taskgroup()

edu_edfi_airflow/providers/earthbeam/operators.py

@@ -118,7 +118,7 @@ class LightbeamOperator(BashOperator):
 
     """
     template_fields = ('data_dir', 'state_dir', 'arguments', 'bash_command', 'env',)
-    valid_commands = ('validate', 'send', 'validate+send')
+    valid_commands = ('validate', 'send', 'validate+send', 'fetch', 'delete')
 
     def __init__(self,
         *,
@@ -140,6 +140,9 @@ def __init__(self,
 
         older_than: Optional[str] = None,
         newer_than: Optional[str] = None,
+        query: Optional[str] = None,
+        keep_keys: Optional[str] = None,
+        drop_keys: Optional[str] = None,
         resend_status_codes: Optional[Union[str, Iterable[str]]] = None,
 
         **kwargs
@@ -191,6 +194,12 @@ def __init__(self,
             self.arguments['--older-than'] = older_than
         if newer_than:
             self.arguments['--newer-than'] = newer_than
+        if query:
+            self.arguments['--query'] = query
+        if keep_keys:
+            self.arguments['--keep-keys'] = keep_keys
+        if drop_keys:
+            self.arguments['--drop-keys'] = drop_keys
 
         ### Environment variables
         # Pass required `data_dir`
@@ -234,8 +243,9 @@ def execute(self, conf=None, **context) -> str:
             logging.info("Parameter `force` provided in context will overwrite defined operator argument.")
             self.arguments['--force'] = ""
 
-        # Create state_dir if not already defined in filespace
+        # Create state_dir and data_dir if not already defined in filespace
         os.makedirs(self.state_dir, exist_ok=True)
+        os.makedirs(self.data_dir, exist_ok=True)
 
         # Format values before adding to the bash command.
         cli_arguments = []