Skip to content
This repository was archived by the owner on May 22, 2025. It is now read-only.
This repository was archived by the owner on May 22, 2025. It is now read-only.

Positional arg to dataframe function throws KeyError #1934

Open
@jeremybeard

Description

@jeremybeard

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions