astronomer/astro-sdk

Positional arg to dataframe function throws KeyError

jeremybeard opened this issue · 0 comments

Describe the bug
When passing a positional argument to an @aql.dataframe decorated function, the SDK fails to start the task because of a KeyError on the parameter.

From looking at the code here my suspicion is that it is assuming that all arguments are in op_kwargs when positional arguments would be in op_args.

Version

  • Astro SDK: 1.5.3
  • Astro runtime: 8.1.0

To Reproduce
Steps to reproduce the behavior:

  1. Manually trigger the DAG:
from airflow.decorators import dag
from astro import sql as aql
import pendulum

@aql.dataframe(task_id="cell_1")
def cell_1_func(a, b, /, c, d, *vararg, e, f, **kwargs):
    print(f'positional only args: {a}, {b}')
    print(f'args: {c}, {d}')
    print(f'keyword only args: {e}, {f}')
    print(f'vararg: {vararg}')
    print(f'kwargs: {kwargs}')
    return

@dag(
    schedule=None,
    start_date=pendulum.from_format("2023-01-01", "YYYY-MM-DD"),
    catchup=False,
)
def my_dag():
    cell_1 = cell_1_func(13, 12, 11, 10, 9, 8, 7, 6, 5, e=4, f=3, g=2, h=1)

dag_obj = my_dag()
  1. See error in failed task log:
[2023-05-10, 17:30:18 UTC] {taskinstance.py:1847} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/astro/sql/operators/dataframe.py", line 149, in execute
    first_table = find_first_table(
  File "/usr/local/lib/python3.10/site-packages/astro/utils/table.py", line 105, in find_first_table
    first_table = _find_first_table_from_op_kwargs(
  File "/usr/local/lib/python3.10/site-packages/astro/utils/table.py", line 50, in _find_first_table_from_op_kwargs
    kwargs = [
  File "/usr/local/lib/python3.10/site-packages/astro/utils/table.py", line 52, in <listcomp>
    if isinstance(op_kwargs[kwarg.name], XComArg)
KeyError: 'a'

Expected behavior
The task to run.