fix: empty openlineage dataset name for AthenaExtractor (#39677)

Signed-off-by: Kacper Muda <mudakacper@gmail.com>
apache · May 17, 2024 · 8d1bd34 · 8d1bd34
1 parent 74c3fb3
commit 8d1bd34
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/airflow/providers/amazon/aws/operators/athena.py b/airflow/providers/amazon/aws/operators/athena.py
@@ -266,7 +266,7 @@ def get_openlineage_facets_on_start(self) -> OperatorLineage:
 
         if self.output_location:
             parsed = urlparse(self.output_location)
-            outputs.append(Dataset(namespace=f"{parsed.scheme}://{parsed.netloc}", name=parsed.path))
+            outputs.append(Dataset(namespace=f"{parsed.scheme}://{parsed.netloc}", name=parsed.path or "/"))
 
         return OperatorLineage(job_facets=job_facets, run_facets=run_facets, inputs=inputs, outputs=outputs)
 

diff --git a/tests/providers/amazon/aws/operators/test_athena.py b/tests/providers/amazon/aws/operators/test_athena.py
@@ -279,7 +279,7 @@ def mock_get_table_metadata(CatalogName, DatabaseName, TableName):
             task_id="test_athena_openlineage",
             query="INSERT INTO TEST_TABLE SELECT CUSTOMER_EMAIL FROM DISCOUNTS",
             database="TEST_DATABASE",
-            output_location="s3://test_s3_bucket/",
+            output_location="s3://test_s3_bucket",
             client_request_token="eac427d0-1c6d-4dfb-96aa-2835d3ac6595",
             sleep_time=0,
             max_polling_attempts=3,