fix(airflow): fix AthenaOperator extraction

[steffengr]

The GenericSqlExtractor which is currently by the DataHub Airflow plugin to extract lineage information does not properly support the AthenaOperator and crashes with "AttributeError: 'AthenaOperator' object has no attribute 'sql'". This patch introduces a AthenaOperatorExtractor following the BigQueryInsertJobOperatorExtractor example to fix support for the AthenaOperator.

Fixes #11160

Checklist

• The PR conforms to DataHub's Contributing Guideline (particularly Commit Message Format)
• Links to related issues (if applicable)
• Tests for the changes have been added/updated (if applicable)
• Docs related to the changes have been added/updated (if applicable). If a new feature has been added a Usage Guide has been added for the same.
• For any breaking change/potential downtime/deprecation/big changes an entry has been made in Updating DataHub

[hsheth2]

The code looks reasonable - can we add a test for this?

[steffengr]

@hsheth2 Thank for looking at this! I don't see related tests in this module. Could you point me at an example? I don't have much experience with this code base.

[hsheth2]

@steffengr we have an integration test (https://github.com/datahub-project/datahub/blob/master/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py) that runs a number of DAGs (e.g. https://github.com/datahub-project/datahub/blob/master/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py)

That'd probably be the easiest way to do it. Otherwise, we could also do a more targeted unit test in https://github.com/datahub-project/datahub/blob/master/metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py

[steffengr]

@hsheth2 I added a test. I had to set "AIRFLOW__API__AUTH_BACKENDS": "airflow.providers.fab.auth_manager.api.auth.backend.basic_auth", to get these tests to work. Probably some change in Airflow.

[steffengr]

I'm not sure what causes these test failures. I'm also getting them on master.

[hsheth2]

Looks like CI is still failing. I see this in the logs (e.g. here https://github.com/datahub-project/datahub/actions/runs/12052943048/job/33607437487?pr=11857)

Traceback (most recent call last):
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/bin/airflow", line 8, in <module>
    sys.exit(main())
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/__main__.py", line 57, in main
    args.func(args)
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/cli/cli_config.py", line 49, in command
    return func(*args, **kwargs)
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/utils/cli.py", line 114, in wrapper
    return f(*args, **kwargs)
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/utils/providers_configuration_loader.py", line 55, in wrapped_function
    return func(*args, **kwargs)
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/cli/commands/dag_command.py", line 160, in dag_trigger
    message = api_client.trigger_dag(
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/api/client/local_client.py", line 34, in trigger_dag
    dag_run = trigger_dag.trigger_dag(
  File "/home/runner/work/datahub/datahub/metadata-ingestion-modules/airflow-plugin/venv/lib/python3.10/site-packages/airflow/api/common/trigger_dag.py", line 119, in trigger_dag
    raise DagNotFound(f"Dag id {dag_id} not found in DagModel")
airflow.exceptions.DagNotFound: Dag id athena_operator not found in DagModel

I also created #11981 to help make this a bit easier to debug, in case that's useful.

[steffengr]

@hsheth2 I'm lost here. I suspected there might be a race condition between triggering the dag and loading it as I get this error randomly on any of the test dags on my local machine. However, calling the rest API to wait for the dag to be available doesn't seem to help on some of the Airflow versions. Any idea what could cause the dag to not load on some of the Airflow versions?

[steffengr]

@hsheth2 Looks like changing my REST call to check if the dag exists by removing the query parameter that didn't exist in old versions of Airflow did the trick but now I'm getting another strange error regarding the golden file not existing on only one particular version of Airflow (2.4.3). Does this version require anything special?

ignore_paths_v2 = (), ignore_order = True
    def assert_metadata_files_equal(
        output_path: Union[str, os.PathLike],
        golden_path: Union[str, os.PathLike],
        update_golden: bool,
        copy_output: bool,
        ignore_paths: Sequence[str] = (),
        ignore_paths_v2: Sequence[str] = (),
        ignore_order: bool = True,
    ) -> None:
        golden_exists = os.path.isfile(golden_path)
    
        if copy_output:
            shutil.copyfile(str(output_path), str(golden_path) + ".output")
            logger.info(f"Copied output file to {golden_path}.output")
    
        if not update_golden and not golden_exists:
>           raise FileNotFoundError(
                "Golden file does not exist. Please run with the --update-golden-files option to create."
            )
E           FileNotFoundError: Golden file does not exist. Please run with the --update-golden-files option to create.
../../metadata-ingestion/src/datahub/testing/compare_metadata_json.py:56: FileNotFoundError

=========================== short test summary info ============================
FAILED tests/integration/test_plugin.py::test_airflow_plugin[v2_athena_operator_no_dag_listener] - FileNotFoundError: Golden file does not exist. Please run with the --update-golden-files option to create.
======= 1 failed, 19 passed, 3 skipped, 21 warnings in 347.02s (0:05:47) =======

[hsheth2]

@steffengr with airflow 2.4, we support a limited form of the plugin, since certain features weren't available until later versions.

We run tests for that as well, which generates the "no_dag_listener" variants of the golden files. That file golden can only be generated on airflow 2.4.

I believe something like this would generate that missing file.

tox -e py310-airflow24 -- 'tests/integration/test_plugin.py::test_airflow_plugin[v2_athena_operator_no_dag_listener]' --update-golden-files

[steffengr]

@hsheth2 Thanks! That did the trick. Is there any chance to get this into a 0.14 patch release?

[hsheth2]

@steffengr this will likely go into 0.15.0. We will definitely at least cut an rc this week, and might also cut the full release