Positional arg to dataframe function throws KeyError
jeremybeard opened this issue · 0 comments
jeremybeard commented
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:
- 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()
- 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.