diff --git a/README.md b/README.md index 133290a6..15b5a6c9 100644 --- a/README.md +++ b/README.md @@ -82,6 +82,7 @@ A dbt profile can be configured to run against AWS Athena using the following co | schema | Specify the schema (Athena database) to build models into (lowercase **only**) | Required | `dbt` | | database | Specify the database (Data catalog) to build models into (lowercase **only**) | Required | `awsdatacatalog` | | poll_interval | Interval in seconds to use for polling the status of query results in Athena | Optional | `5` | +| debug_query_state | Flag if debug message with Athena query state is needed | Optional | `false` | | aws_access_key_id | Access key ID of the user performing requests. | Optional | `AKIAIOSFODNN7EXAMPLE` | | aws_secret_access_key | Secret access key of the user performing requests | Optional | `wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY` | | aws_profile_name | Profile to use from your AWS shared credentials file. | Optional | `my-profile` | diff --git a/dbt/adapters/athena/connections.py b/dbt/adapters/athena/connections.py index 8645f8bc..89363d92 100644 --- a/dbt/adapters/athena/connections.py +++ b/dbt/adapters/athena/connections.py @@ -54,6 +54,7 @@ class AthenaCredentials(Credentials): aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None poll_interval: float = 1.0 + debug_query_state: bool = False _ALIASES = {"catalog": "database"} num_retries: Optional[int] = 5 s3_data_dir: Optional[str] = None @@ -81,7 +82,7 @@ def _connection_keys(self) -> Tuple[str, ...]: "endpoint_url", "s3_data_dir", "s3_data_naming", - "lf_tags", + "debug_query_state", ) @@ -122,7 +123,8 @@ def __poll(self, query_id: str) -> AthenaQueryExecution: ]: return query_execution else: - logger.debug(f"Query state is: {query_execution.state}. Sleeping for {self._poll_interval}...") + if self.connection.cursor_kwargs.get("debug_query_state", False): + logger.debug(f"Query state is: {query_execution.state}. Sleeping for {self._poll_interval}...") time.sleep(self._poll_interval) def execute( # type: ignore @@ -215,6 +217,7 @@ def open(cls, connection: Connection) -> Connection: schema_name=creds.schema, work_group=creds.work_group, cursor_class=AthenaCursor, + cursor_kwargs={"debug_query_state": creds.debug_query_state}, formatter=AthenaParameterFormatter(), poll_interval=creds.poll_interval, session=get_boto3_session(connection), diff --git a/dev-requirements.txt b/dev-requirements.txt index 41e3249f..1070e5fc 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -4,8 +4,8 @@ dbt-tests-adapter~=1.6.1 flake8~=6.1 Flake8-pyproject~=1.2 isort~=5.11 -moto~=4.2.0 -pre-commit~=2.21 +moto~=4.2.2 +pre-commit~=3.4 pyparsing~=3.1.1 pytest~=7.4 pytest-cov~=4.1