Fix to make expand task params as optional

[ifulia]

Make expand task params as optional, so not all tasks need to provide it. Currently, if there is no expand field in tasks YAML, it fails to generate the DAG from YAML. Also using expand creates a task mapping in airflow which doesn't not work when you want to use a branch operator to skip certain tasks based on the condition. Fix it here by making it optional.

[matveykortsev]

Hey @ifulia, I dont think its actually fix.
Its already possible to make it optional:
https://github.com/ajbosco/dag-factory/blob/aad7600f3df56f6c5cc2c0aa304816f43eb29a32/dagfactory/dagbuilder.py#L579-L589
These checks are made in dag builder, before get_expand_partial_kwargs which you edited.
Here is examples with expand and without expand which is currently in my production:

dag_name:
  description: ""
  schedule_interval: "50 9 * * *"
  template_searchpath: ['/opt/airflow/dags/folder']
  tasks:
    task_1:
      operator: operator.Operator
      conn_id: 'conn_id'
    task_2:
      operator: airflow.operators.python_operator.PythonOperator
      python_callable_name: python_callable
      python_callable_file: /opt/airflow/plugins/utils.py
      dependencies: [task_1]
    task_3:
      operator: airflow.operators.python_operator.PythonOperator
      python_callable_name: prometheus_task
      python_callable_file: /opt/airflow/plugins/utils.py
      partial:
        op_kwargs:
          conn_id: "conn_id"
          prometheus_url: "{{ conn.PROMETHEUS }}"
          pg_host: "{{ conn.PG.host }}"
          pg_schema: "{{ conn.PG.schema }}"
          pg_login: "{{ conn.PG.login }}"
          pg_password: "{{ conn.PG.password }}"
      expand:
        op_args:
          task_2.output
      dependencies:
        - task_2

dag_name:
  template_searchpath: ['/opt/airflow/dags/sql']
  schedule_interval: "0 5 * * MON"
  tasks:
    task_1:
      sql: 'insert_editor_event_week.sql'
      operator: path_to_sql_operator.Operator
      sql_truncate: 'alter.sql'
      conn_id: 'conn_id'
      params:
        table: 'table_name'
        delete_condition: 'timestamp'

[ifulia]

Hey @ifulia, I dont think its actually fix. Its already possible to make it optional:

https://github.com/ajbosco/dag-factory/blob/aad7600f3df56f6c5cc2c0aa304816f43eb29a32/dagfactory/dagbuilder.py#L579-L589

These checks are made in dag builder, before get_expand_partial_kwargs which you edited.
Here is examples with expand and without expand which is currently in my production:

dag_name:
  description: ""
  schedule_interval: "50 9 * * *"
  template_searchpath: ['/opt/airflow/dags/folder']
  tasks:
    task_1:
      operator: operator.Operator
      conn_id: 'conn_id'
    task_2:
      operator: airflow.operators.python_operator.PythonOperator
      python_callable_name: python_callable
      python_callable_file: /opt/airflow/plugins/utils.py
      dependencies: [task_1]
    task_3:
      operator: airflow.operators.python_operator.PythonOperator
      python_callable_name: prometheus_task
      python_callable_file: /opt/airflow/plugins/utils.py
      partial:
        op_kwargs:
          conn_id: "conn_id"
          prometheus_url: "{{ conn.PROMETHEUS }}"
          pg_host: "{{ conn.PG.host }}"
          pg_schema: "{{ conn.PG.schema }}"
          pg_login: "{{ conn.PG.login }}"
          pg_password: "{{ conn.PG.password }}"
      expand:
        op_args:
          task_2.output
      dependencies:
        - task_2

dag_name:
  template_searchpath: ['/opt/airflow/dags/sql']
  schedule_interval: "0 5 * * MON"
  tasks:
    task_1:
      sql: 'insert_editor_event_week.sql'
      operator: path_to_sql_operator.Operator
      sql_truncate: 'alter.sql'
      conn_id: 'conn_id'
      params:
        table: 'table_name'
        delete_condition: 'timestamp'

@matveykortsev
If you don't use partial and expand then it works. Problem comes when you just use partial.

Here is what I am using and getting error when creating dag.

dag_name:
  default_args:
    owner: 'example_owner'
    retries: 1
    start_date: 2018-01-01
    retry_delay_sec: 300
  description: ""
  schedule_interval: "50 9 * * *"
  tasks:
    task_1:
      operator: airflow.operators.python.PythonOperator
      python_callable_name: start
      python_callable_file: /opt/airflow/dags/create_env.py
    task_2:
      operator: airflow.operators.python.PythonOperator
      python_callable_name: join
      python_callable_file: /opt/airflow/dags/create_env.py
      partial:
        op_kwargs:
          env_name: "test"
      dependencies: [task_1]

The error I get is

Traceback (most recent call last):
  File "/opt/homebrew/lib/python3.11/site-packages/dagfactory/dagbuilder.py", line 589, in make_task
    ) = utils.get_expand_partial_kwargs(task_params)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/dagfactory/utils.py", line 235, in get_expand_partial_kwargs
    for expand_key, expand_value in task_params["expand"].items():
                                    ~~~~~~~~~~~^^^^^^^^^^
KeyError: 'expand'

[ifulia]

@matveykortsev @ajbosco Could you please help merge this PR.

[ifulia]

@ajbosco @matveykortsev Could you please help review and merge this. This is blocking our use-case where we want to use only partial params. Let me know if you have any concerns.

[matveykortsev]

@ajbosco @matveykortsev Could you please help review and merge this. This is blocking our use-case where we want to use only partial params. Let me know if you have any concerns.

I still don’t understand of use case with partial but without expand. You can just pass params directly to operators then

[ifulia]

I still don’t understand of use case with partial but without expand. You can just pass params directly to operators then

Ah, I didnt realize that you can pass it with params. I thought the only way to pass params is through partial and/or expand. This solves the problem. Thanks @matveykortsev
Closing this PR