From 0f037432211272941a3368e2c5db8400a166f0eb Mon Sep 17 00:00:00 2001 From: Jakub Dardzinski Date: Fri, 20 Oct 2023 22:54:18 +0200 Subject: [PATCH 1/3] Add OpenLineage support for AthenaOperator. Based on: https://github.com/OpenLineage/OpenLineage/pull/1328. Signed-off-by: Jakub Dardzinski --- .../providers/amazon/aws/operators/athena.py | 133 ++++++++++++++++++ .../amazon/aws/operators/athena_metadata.json | 72 ++++++++++ .../amazon/aws/operators/test_athena.py | 120 +++++++++++++++- 3 files changed, 324 insertions(+), 1 deletion(-) create mode 100644 tests/providers/amazon/aws/operators/athena_metadata.json diff --git a/airflow/providers/amazon/aws/operators/athena.py b/airflow/providers/amazon/aws/operators/athena.py index 1f9e527717f61..d211f61aaff73 100644 --- a/airflow/providers/amazon/aws/operators/athena.py +++ b/airflow/providers/amazon/aws/operators/athena.py @@ -18,6 +18,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, Sequence +from urllib.parse import urlparse from airflow.configuration import conf from airflow.exceptions import AirflowException @@ -27,6 +28,9 @@ from airflow.providers.amazon.aws.utils.mixins import aws_template_fields if TYPE_CHECKING: + from openlineage.client.facet import BaseFacet + from openlineage.client.run import Dataset + from airflow.utils.context import Context @@ -187,3 +191,132 @@ def on_kill(self) -> None: "Polling Athena for query with id %s to reach final state", self.query_execution_id ) self.hook.poll_query_status(self.query_execution_id, sleep_time=self.sleep_time) + + def get_openlineage_facets_on_start(self): + from openlineage.client.facet import ExtractionError, ExtractionErrorRunFacet, SqlJobFacet + from openlineage.client.run import Dataset + + from airflow.providers.openlineage.extractors.base import OperatorLineage + from airflow.providers.openlineage.sqlparser import SQLParser + + sql_parser = SQLParser(dialect="generic") + + job_facets: dict[str, BaseFacet] = {"sql": SqlJobFacet(query=sql_parser.normalize_sql(self.query))} + parse_result = sql_parser.parse(sql=self.query) + + if not parse_result: + return OperatorLineage(job_facets=job_facets) + + run_facets: dict[str, BaseFacet] = {} + if parse_result.errors: + run_facets["extractionError"] = ExtractionErrorRunFacet( + totalTasks=len(self.query) if isinstance(self.query, list) else 1, + failedTasks=len(parse_result.errors), + errors=[ + ExtractionError( + errorMessage=error.message, + stackTrace=None, + task=error.origin_statement, + taskNumber=error.index, + ) + for error in parse_result.errors + ], + ) + + inputs: list[Dataset] = list( + filter( + None, + [ + self.get_openlineage_dataset(table.schema or self.database, table.name) + for table in parse_result.in_tables + ], + ) + ) + + # Athena can output query result to a new table with CTAS query. + # cf. https://docs.aws.amazon.com/athena/latest/ug/ctas.html + outputs: list[Dataset] = list( + filter( + None, + [ + self.get_openlineage_dataset(table.schema or self.database, table.name) + for table in parse_result.out_tables + ], + ) + ) + + # In addition to CTAS query, it's also possible to specify output location on S3 + # with a mandatory parameter, OutputLocation in ResultConfiguration. + # cf. https://docs.aws.amazon.com/athena/latest/APIReference/API_ResultConfiguration.html#athena-Type-ResultConfiguration-OutputLocation # noqa: E501 + # + # Depending on the query type and the external_location property in the CTAS query, + # its behavior changes as follows: + # + # * Normal SELECT statement + # -> The result is put into output_location as files rather than a table. + # + # * CTAS statement without external_location (`CREATE TABLE ... AS SELECT ...`) + # -> The result is put into output_location as a table, + # that is, both metadata files and data files are in there. + # + # * CTAS statement with external_location + # (`CREATE TABLE ... WITH (external_location='s3://bucket/key') AS SELECT ...`) + # -> The result is output as a table, but metadata and data files are + # separated into output_location and external_location respectively. + # + # For the last case, output_location may be unnecessary as OL's output information, + # but we keep it as of now since it may be useful for some purpose. + output_location = self.output_location + parsed = urlparse(output_location) + outputs.append(Dataset(namespace=f"{parsed.scheme}://{parsed.netloc}", name=parsed.path)) + + return OperatorLineage(job_facets=job_facets, run_facets=run_facets, inputs=inputs, outputs=outputs) + + def get_openlineage_dataset(self, database, table) -> Dataset | None: + from openlineage.client.facet import ( + SchemaDatasetFacet, + SchemaField, + SymlinksDatasetFacet, + SymlinksDatasetFacetIdentifiers, + ) + from openlineage.client.run import Dataset + + # Currently, AthenaOperator and AthenaHook don't have a functionality to specify catalog, + # and it seems to implicitly assume that the default catalog (AwsDataCatalog) is target. + CATALOG_NAME = "AwsDataCatalog" + + client = self.hook.get_conn() + try: + table_metadata = client.get_table_metadata( + CatalogName=CATALOG_NAME, DatabaseName=database, TableName=table + ) + + # Dataset has also its' physical location which we can add in symlink facet. + s3_location = table_metadata["TableMetadata"]["Parameters"]["location"] + parsed_path = urlparse(s3_location) + facets: dict[str, BaseFacet] = { + "symlinks": SymlinksDatasetFacet( + identifiers=[ + SymlinksDatasetFacetIdentifiers( + namespace=f"{parsed_path.scheme}://{parsed_path.netloc}", # type: ignore + name=str(parsed_path.path), + type="TABLE", + ) + ] + ) + } + fields = [ + SchemaField(name=column["Name"], type=column["Type"], description=column["Comment"]) + for column in table_metadata["TableMetadata"]["Columns"] + ] + if fields: + facets["schema"] = SchemaDatasetFacet(fields=fields) + return Dataset( + namespace=f"awsathena://athena.{self.hook.region_name}.amazonaws.com", + name=".".join(filter(None, (CATALOG_NAME, database, table))), + facets=facets, + ) + + except Exception as e: + self.log.error("Cannot retrieve table metadata from Athena.Client. %s", e) + return None diff --git a/tests/providers/amazon/aws/operators/athena_metadata.json b/tests/providers/amazon/aws/operators/athena_metadata.json new file mode 100644 index 0000000000000..f13b1241742ab --- /dev/null +++ b/tests/providers/amazon/aws/operators/athena_metadata.json @@ -0,0 +1,72 @@ +{ + "DISCOUNTS": { + "TableMetadata": { + "Name": "DISCOUNTS", + "CreateTime": 1593559968.0, + "LastAccessTime": 0.0, + "TableType": "EXTERNAL_TABLE", + "Columns": [ + { + "Name": "ID", + "Type": "int", + "Comment": "from deserializer" + }, + { + "Name": "AMOUNT_OFF", + "Type": "int", + "Comment": "from deserializer" + }, + { + "Name": "CUSTOMER_EMAIL", + "Type": "varchar", + "Comment": "from deserializer" + }, + { + "Name": "STARTS_ON", + "Type": "timestamp", + "Comment": "from deserializer" + }, + { + "Name": "ENDS_ON", + "Type": "timestamp", + "Comment": "from deserializer" + } + ], + "PartitionKeys": [], + "Parameters": { + "EXTERNAL": "TRUE", + "inputformat": "com.esri.json.hadoop.EnclosedJsonInputFormat", + "location": "s3://bucket/discount/data/path/", + "outputformat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat", + "serde.param.serialization.format": "1", + "serde.serialization.lib": "com.esri.hadoop.hive.serde.JsonSerde", + "transient_lastDdlTime": "1593559968" + } + } + }, + "TEST_TABLE": { + "TableMetadata": { + "Name": "TEST_TABLE", + "CreateTime": 1593559968.0, + "LastAccessTime": 0.0, + "TableType": "EXTERNAL_TABLE", + "Columns": [ + { + "Name": "column", + "Type": "string", + "Comment": "from deserializer" + } + ], + "PartitionKeys": [], + "Parameters": { + "EXTERNAL": "TRUE", + "inputformat": "com.esri.json.hadoop.EnclosedJsonInputFormat", + "location": "s3://bucket/data/test_table/data/path", + "outputformat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat", + "serde.param.serialization.format": "1", + "serde.serialization.lib": "com.esri.hadoop.hive.serde.JsonSerde", + "transient_lastDdlTime": "1593559968" + } + } + } + } diff --git a/tests/providers/amazon/aws/operators/test_athena.py b/tests/providers/amazon/aws/operators/test_athena.py index 497e2c212774d..5820827b05714 100644 --- a/tests/providers/amazon/aws/operators/test_athena.py +++ b/tests/providers/amazon/aws/operators/test_athena.py @@ -16,15 +16,25 @@ # under the License. from __future__ import annotations +import json from unittest import mock import pytest +from openlineage.client.facet import ( + SchemaDatasetFacet, + SchemaField, + SqlJobFacet, + SymlinksDatasetFacet, + SymlinksDatasetFacetIdentifiers, +) +from openlineage.client.run import Dataset from airflow.exceptions import TaskDeferred from airflow.models import DAG, DagRun, TaskInstance from airflow.providers.amazon.aws.hooks.athena import AthenaHook from airflow.providers.amazon.aws.operators.athena import AthenaOperator from airflow.providers.amazon.aws.triggers.athena import AthenaTrigger +from airflow.providers.openlineage.extractors import OperatorLineage from airflow.utils import timezone from airflow.utils.timezone import datetime @@ -150,7 +160,11 @@ def test_hook_run_big_success_query(self, mock_conn, mock_run_query, mock_check_ @mock.patch.object(AthenaHook, "run_query", return_value=ATHENA_QUERY_ID) @mock.patch.object(AthenaHook, "get_conn") def test_hook_run_failure_query( - self, mock_conn, mock_run_query, mock_check_query_status, mock_get_state_change_reason + self, + mock_conn, + mock_run_query, + mock_check_query_status, + mock_get_state_change_reason, ): with pytest.raises(Exception): self.athena.execute({}) @@ -226,3 +240,107 @@ def test_is_deferred(self, mock_run_query): self.athena.execute(None) assert isinstance(deferred.value.trigger, AthenaTrigger) + + @mock.patch.object(AthenaHook, "region_name", new_callable=mock.PropertyMock) + @mock.patch.object(AthenaHook, "get_conn") + def test_operator_openlineage_data(self, mock_conn, mock_region_name): + mock_region_name.return_value = "eu-west-1" + + def mock_get_table_metadata(CatalogName, DatabaseName, TableName): + with open("tests/providers/amazon/aws/operators/athena_metadata.json") as f: + return json.load(f)[TableName] + + mock_conn.return_value.get_table_metadata = mock_get_table_metadata + + op = AthenaOperator( + task_id="test_athena_openlineage", + query="INSERT INTO TEST_TABLE SELECT CUSTOMER_EMAIL FROM DISCOUNTS", + database="TEST_DATABASE", + output_location="s3://test_s3_bucket/", + client_request_token="eac427d0-1c6d-4dfb-96aa-2835d3ac6595", + sleep_time=0, + max_polling_attempts=3, + dag=self.dag, + ) + + expected_lineage = OperatorLineage( + inputs=[ + Dataset( + namespace="awsathena://athena.eu-west-1.amazonaws.com", + name="AwsDataCatalog.TEST_DATABASE.DISCOUNTS", + facets={ + "symlinks": SymlinksDatasetFacet( + identifiers=[ + SymlinksDatasetFacetIdentifiers( + namespace="s3://bucket", + name="/discount/data/path/", + type="TABLE", + ) + ], + ), + "schema": SchemaDatasetFacet( + fields=[ + SchemaField( + name="ID", + type="int", + description="from deserializer", + ), + SchemaField( + name="AMOUNT_OFF", + type="int", + description="from deserializer", + ), + SchemaField( + name="CUSTOMER_EMAIL", + type="varchar", + description="from deserializer", + ), + SchemaField( + name="STARTS_ON", + type="timestamp", + description="from deserializer", + ), + SchemaField( + name="ENDS_ON", + type="timestamp", + description="from deserializer", + ), + ], + ), + }, + ) + ], + outputs=[ + Dataset( + namespace="awsathena://athena.eu-west-1.amazonaws.com", + name="AwsDataCatalog.TEST_DATABASE.TEST_TABLE", + facets={ + "symlinks": SymlinksDatasetFacet( + identifiers=[ + SymlinksDatasetFacetIdentifiers( + namespace="s3://bucket", + name="/data/test_table/data/path", + type="TABLE", + ) + ], + ), + "schema": SchemaDatasetFacet( + fields=[ + SchemaField( + name="column", + type="string", + description="from deserializer", + ) + ], + ), + }, + ), + Dataset(namespace="s3://test_s3_bucket", name="/"), + ], + job_facets={ + "sql": SqlJobFacet( + query="INSERT INTO TEST_TABLE SELECT CUSTOMER_EMAIL FROM DISCOUNTS", + ) + }, + ) + assert op.get_openlineage_facets_on_start() == expected_lineage From 9a449283dd582d05525523249e6e73d5f2f4162f Mon Sep 17 00:00:00 2001 From: Jakub Dardzinski Date: Wed, 25 Oct 2023 01:00:41 +0200 Subject: [PATCH 2/3] Cache `get_query_execution` in AthenaHook. Adjust code to catalog and output_location changes. Signed-off-by: Jakub Dardzinski --- airflow/providers/amazon/aws/hooks/athena.py | 38 ++++++++++++-- .../providers/amazon/aws/operators/athena.py | 50 ++++++------------- .../providers/amazon/aws/hooks/test_athena.py | 17 +++++++ 3 files changed, 66 insertions(+), 39 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/athena.py b/airflow/providers/amazon/aws/hooks/athena.py index 11148bd6b17a5..8f7ad9c36e466 100644 --- a/airflow/providers/amazon/aws/hooks/athena.py +++ b/airflow/providers/amazon/aws/hooks/athena.py @@ -27,6 +27,7 @@ import warnings from typing import TYPE_CHECKING, Any +from airflow.compat.functools import cache from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.utils.waiter_with_logging import wait @@ -120,7 +121,26 @@ def run_query( self.log.info("Query execution id: %s", query_execution_id) return query_execution_id - def check_query_status(self, query_execution_id: str) -> str | None: + def get_query_execution(self, query_execution_id: str, use_cached_response: bool = False) -> dict: + """Get information about a single execution of a query. + + .. seealso:: + - :external+boto3:py:meth:`Athena.Client.get_query_execution` + + :param query_execution_id: Id of submitted athena query + :param use_cached_response: If True, use execution information cache + """ + # second check if to satisfy mypy + if use_cached_response or not hasattr(self._get_query_execution, "__wrapped__"): + return self._get_query_execution(query_execution_id=query_execution_id) + return self._get_query_execution.__wrapped__(self, query_execution_id=query_execution_id) + + @cache + def _get_query_execution(self, query_execution_id: str): + response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id) + return response + + def check_query_status(self, query_execution_id: str, use_cached_response: bool = False) -> str | None: """Fetch the state of a submitted query. .. seealso:: @@ -130,7 +150,9 @@ def check_query_status(self, query_execution_id: str) -> str | None: :return: One of valid query states, or *None* if the response is malformed. """ - response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id) + response = self.get_query_execution( + query_execution_id=query_execution_id, use_cached_response=use_cached_response + ) state = None try: state = response["QueryExecution"]["Status"]["State"] @@ -143,7 +165,9 @@ def check_query_status(self, query_execution_id: str) -> str | None: # The error is being absorbed to implement retries. return state - def get_state_change_reason(self, query_execution_id: str) -> str | None: + def get_state_change_reason( + self, query_execution_id: str, use_cached_response: bool = False + ) -> str | None: """ Fetch the reason for a state change (e.g. error message). Returns None or reason string. @@ -152,7 +176,9 @@ def get_state_change_reason(self, query_execution_id: str) -> str | None: :param query_execution_id: Id of submitted athena query """ - response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id) + response = self.get_query_execution( + query_execution_id=query_execution_id, use_cached_response=use_cached_response + ) reason = None try: reason = response["QueryExecution"]["Status"]["StateChangeReason"] @@ -277,7 +303,9 @@ def get_output_location(self, query_execution_id: str) -> str: """ output_location = None if query_execution_id: - response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id) + response = self.get_query_execution( + query_execution_id=query_execution_id, use_cached_response=True + ) if response: try: diff --git a/airflow/providers/amazon/aws/operators/athena.py b/airflow/providers/amazon/aws/operators/athena.py index d211f61aaff73..3bc907227d603 100644 --- a/airflow/providers/amazon/aws/operators/athena.py +++ b/airflow/providers/amazon/aws/operators/athena.py @@ -31,6 +31,7 @@ from openlineage.client.facet import BaseFacet from openlineage.client.run import Dataset + from airflow.providers.openlineage.extractors.base import OperatorLineage from airflow.utils.context import Context @@ -164,6 +165,9 @@ def execute(self, context: Context) -> str | None: f"query_execution_id is {self.query_execution_id}." ) + # Save output location from API response for later use in OpenLineage. + self.output_location = self.hook.get_output_location(self.query_execution_id) + return self.query_execution_id def execute_complete(self, context, event=None): @@ -192,7 +196,12 @@ def on_kill(self) -> None: ) self.hook.poll_query_status(self.query_execution_id, sleep_time=self.sleep_time) - def get_openlineage_facets_on_start(self): + def get_openlineage_facets_on_start(self) -> OperatorLineage: + """Retrieve OpenLineage data by parsing SQL queries and enriching them with Athena API. + + In addition to CTAS query, query and calculation results are stored in S3 location. + For that reason additional output is attached with this location. + """ from openlineage.client.facet import ExtractionError, ExtractionErrorRunFacet, SqlJobFacet from openlineage.client.run import Dataset @@ -233,8 +242,6 @@ def get_openlineage_facets_on_start(self): ) ) - # Athena can output query result to a new table with CTAS query. - # cf. https://docs.aws.amazon.com/athena/latest/ug/ctas.html outputs: list[Dataset] = list( filter( None, @@ -245,30 +252,9 @@ def get_openlineage_facets_on_start(self): ) ) - # In addition to CTAS query, it's also possible to specify output location on S3 - # with a mandatory parameter, OutputLocation in ResultConfiguration. - # cf. https://docs.aws.amazon.com/athena/latest/APIReference/API_ResultConfiguration.html#athena-Type-ResultConfiguration-OutputLocation # noqa: E501 - # - # Depending on the query type and the external_location property in the CTAS query, - # its behavior changes as follows: - # - # * Normal SELECT statement - # -> The result is put into output_location as files rather than a table. - # - # * CTAS statement without external_location (`CREATE TABLE ... AS SELECT ...`) - # -> The result is put into output_location as a table, - # that is, both metadata files and data files are in there. - # - # * CTAS statement with external_location - # (`CREATE TABLE ... WITH (external_location='s3://bucket/key') AS SELECT ...`) - # -> The result is output as a table, but metadata and data files are - # separated into output_location and external_location respectively. - # - # For the last case, output_location may be unnecessary as OL's output information, - # but we keep it as of now since it may be useful for some purpose. - output_location = self.output_location - parsed = urlparse(output_location) - outputs.append(Dataset(namespace=f"{parsed.scheme}://{parsed.netloc}", name=parsed.path)) + if self.output_location: + parsed = urlparse(self.output_location) + outputs.append(Dataset(namespace=f"{parsed.scheme}://{parsed.netloc}", name=parsed.path)) return OperatorLineage(job_facets=job_facets, run_facets=run_facets, inputs=inputs, outputs=outputs) @@ -281,14 +267,10 @@ def get_openlineage_dataset(self, database, table) -> Dataset | None: ) from openlineage.client.run import Dataset - # Currently, AthenaOperator and AthenaHook don't have a functionality to specify catalog, - # and it seems to implicitly assume that the default catalog (AwsDataCatalog) is target. - CATALOG_NAME = "AwsDataCatalog" - client = self.hook.get_conn() try: table_metadata = client.get_table_metadata( - CatalogName=CATALOG_NAME, DatabaseName=database, TableName=table + CatalogName=self.catalog, DatabaseName=database, TableName=table ) # Dataset has also its' physical location which we can add in symlink facet. @@ -298,7 +280,7 @@ def get_openlineage_dataset(self, database, table) -> Dataset | None: "symlinks": SymlinksDatasetFacet( identifiers=[ SymlinksDatasetFacetIdentifiers( - namespace=f"{parsed_path.scheme}://{parsed_path.netloc}", # type: ignore + namespace=f"{parsed_path.scheme}://{parsed_path.netloc}", name=str(parsed_path.path), type="TABLE", ) @@ -313,7 +295,7 @@ def get_openlineage_dataset(self, database, table) -> Dataset | None: facets["schema"] = SchemaDatasetFacet(fields=fields) return Dataset( namespace=f"awsathena://athena.{self.hook.region_name}.amazonaws.com", - name=".".join(filter(None, (CATALOG_NAME, database, table))), + name=".".join(filter(None, (self.catalog, database, table))), facets=facets, ) diff --git a/tests/providers/amazon/aws/hooks/test_athena.py b/tests/providers/amazon/aws/hooks/test_athena.py index 05ed6e9e307a4..2f0af4b346fa3 100644 --- a/tests/providers/amazon/aws/hooks/test_athena.py +++ b/tests/providers/amazon/aws/hooks/test_athena.py @@ -43,6 +43,7 @@ "QueryExecution": { "QueryExecutionId": MOCK_DATA["query_execution_id"], "ResultConfiguration": {"OutputLocation": "s3://test_bucket/test.csv"}, + "Status": {"StateChangeReason": "Terminated by user."}, } } @@ -50,6 +51,7 @@ class TestAthenaHook: def setup_method(self): self.athena = AthenaHook() + self.athena._get_query_execution.cache_clear() def test_init(self): assert self.athena.aws_conn_id == "aws_default" @@ -195,3 +197,18 @@ def test_hook_get_output_location(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_QUERY_EXECUTION_OUTPUT result = self.athena.get_output_location(query_execution_id=MOCK_DATA["query_execution_id"]) assert result == "s3://test_bucket/test.csv" + + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_get_query_execution_caching(self, mock_conn): + mock_conn.return_value.get_query_execution.return_value = MOCK_QUERY_EXECUTION_OUTPUT + self.athena.get_state_change_reason(query_execution_id=MOCK_DATA["query_execution_id"]) + assert self.athena._get_query_execution.cache_info().misses == 0 + assert self.athena._get_query_execution.cache_info().hits == 0 + # get_output_location uses cache + self.athena.get_output_location(query_execution_id=MOCK_DATA["query_execution_id"]) + assert self.athena._get_query_execution.cache_info().misses == 1 + self.athena.get_state_change_reason( + query_execution_id=MOCK_DATA["query_execution_id"], use_cached_response=False + ) + assert self.athena._get_query_execution.cache_info().hits == 0 + assert self.athena._get_query_execution.cache_info().misses == 1 From 94eb315428cf0ddadd5de191a049e068a34fbf65 Mon Sep 17 00:00:00 2001 From: Jakub Dardzinski Date: Fri, 27 Oct 2023 14:01:40 +0200 Subject: [PATCH 3/3] Change caching implementation. Rename `get_query_execution` to `get_query_info`. Signed-off-by: Jakub Dardzinski --- airflow/providers/amazon/aws/hooks/athena.py | 35 +++++++------------ .../providers/amazon/aws/hooks/test_athena.py | 18 +++++----- 2 files changed, 22 insertions(+), 31 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/athena.py b/airflow/providers/amazon/aws/hooks/athena.py index 8f7ad9c36e466..c8af86ee3a183 100644 --- a/airflow/providers/amazon/aws/hooks/athena.py +++ b/airflow/providers/amazon/aws/hooks/athena.py @@ -27,7 +27,6 @@ import warnings from typing import TYPE_CHECKING, Any -from airflow.compat.functools import cache from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.utils.waiter_with_logging import wait @@ -83,6 +82,7 @@ def __init__( else: self.sleep_time = 30 # previous default value self.log_query = log_query + self.__query_results: dict[str, Any] = {} def run_query( self, @@ -121,26 +121,23 @@ def run_query( self.log.info("Query execution id: %s", query_execution_id) return query_execution_id - def get_query_execution(self, query_execution_id: str, use_cached_response: bool = False) -> dict: + def get_query_info(self, query_execution_id: str, use_cache: bool = False) -> dict: """Get information about a single execution of a query. .. seealso:: - :external+boto3:py:meth:`Athena.Client.get_query_execution` :param query_execution_id: Id of submitted athena query - :param use_cached_response: If True, use execution information cache + :param use_cache: If True, use execution information cache """ - # second check if to satisfy mypy - if use_cached_response or not hasattr(self._get_query_execution, "__wrapped__"): - return self._get_query_execution(query_execution_id=query_execution_id) - return self._get_query_execution.__wrapped__(self, query_execution_id=query_execution_id) - - @cache - def _get_query_execution(self, query_execution_id: str): + if use_cache and query_execution_id in self.__query_results: + return self.__query_results[query_execution_id] response = self.get_conn().get_query_execution(QueryExecutionId=query_execution_id) + if use_cache: + self.__query_results[query_execution_id] = response return response - def check_query_status(self, query_execution_id: str, use_cached_response: bool = False) -> str | None: + def check_query_status(self, query_execution_id: str, use_cache: bool = False) -> str | None: """Fetch the state of a submitted query. .. seealso:: @@ -150,9 +147,7 @@ def check_query_status(self, query_execution_id: str, use_cached_response: bool :return: One of valid query states, or *None* if the response is malformed. """ - response = self.get_query_execution( - query_execution_id=query_execution_id, use_cached_response=use_cached_response - ) + response = self.get_query_info(query_execution_id=query_execution_id, use_cache=use_cache) state = None try: state = response["QueryExecution"]["Status"]["State"] @@ -165,9 +160,7 @@ def check_query_status(self, query_execution_id: str, use_cached_response: bool # The error is being absorbed to implement retries. return state - def get_state_change_reason( - self, query_execution_id: str, use_cached_response: bool = False - ) -> str | None: + def get_state_change_reason(self, query_execution_id: str, use_cache: bool = False) -> str | None: """ Fetch the reason for a state change (e.g. error message). Returns None or reason string. @@ -176,9 +169,7 @@ def get_state_change_reason( :param query_execution_id: Id of submitted athena query """ - response = self.get_query_execution( - query_execution_id=query_execution_id, use_cached_response=use_cached_response - ) + response = self.get_query_info(query_execution_id=query_execution_id, use_cache=use_cache) reason = None try: reason = response["QueryExecution"]["Status"]["StateChangeReason"] @@ -303,9 +294,7 @@ def get_output_location(self, query_execution_id: str) -> str: """ output_location = None if query_execution_id: - response = self.get_query_execution( - query_execution_id=query_execution_id, use_cached_response=True - ) + response = self.get_query_info(query_execution_id=query_execution_id, use_cache=True) if response: try: diff --git a/tests/providers/amazon/aws/hooks/test_athena.py b/tests/providers/amazon/aws/hooks/test_athena.py index 2f0af4b346fa3..8f224f0b2de5b 100644 --- a/tests/providers/amazon/aws/hooks/test_athena.py +++ b/tests/providers/amazon/aws/hooks/test_athena.py @@ -51,7 +51,6 @@ class TestAthenaHook: def setup_method(self): self.athena = AthenaHook() - self.athena._get_query_execution.cache_clear() def test_init(self): assert self.athena.aws_conn_id == "aws_default" @@ -199,16 +198,19 @@ def test_hook_get_output_location(self, mock_conn): assert result == "s3://test_bucket/test.csv" @mock.patch.object(AthenaHook, "get_conn") - def test_hook_get_query_execution_caching(self, mock_conn): + def test_hook_get_query_info_caching(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_QUERY_EXECUTION_OUTPUT self.athena.get_state_change_reason(query_execution_id=MOCK_DATA["query_execution_id"]) - assert self.athena._get_query_execution.cache_info().misses == 0 - assert self.athena._get_query_execution.cache_info().hits == 0 + assert not self.athena._AthenaHook__query_results # get_output_location uses cache self.athena.get_output_location(query_execution_id=MOCK_DATA["query_execution_id"]) - assert self.athena._get_query_execution.cache_info().misses == 1 + assert MOCK_DATA["query_execution_id"] in self.athena._AthenaHook__query_results + mock_conn.return_value.get_query_execution.assert_called_with( + QueryExecutionId=MOCK_DATA["query_execution_id"] + ) self.athena.get_state_change_reason( - query_execution_id=MOCK_DATA["query_execution_id"], use_cached_response=False + query_execution_id=MOCK_DATA["query_execution_id"], use_cache=False + ) + mock_conn.return_value.get_query_execution.assert_called_with( + QueryExecutionId=MOCK_DATA["query_execution_id"] ) - assert self.athena._get_query_execution.cache_info().hits == 0 - assert self.athena._get_query_execution.cache_info().misses == 1