vincbeck commented on code in PR #75997:
URL: https://github.com/apache/airflow/pull/75997#discussion_r4425058866
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/duckdb.py:
##########
@@ -196,6 +200,25 @@ def get_s3_endpoint_url(self) -> str | None:
return explicit
return self.connection_extra.get("endpoint_url") or
self.aws_hook.conn_config.endpoint_url
+ def get_conn(self) -> DuckDBPyConnection:
+ """Return a DuckDB connection, and log the AWS identity once to help
debug S3 ``HTTP 403`` errors."""
+ conn = super().get_conn()
+ if self._aws_identity_looked_up:
+ return conn
+ self._aws_identity_looked_up = True
+ try:
+ if self.credential_strategy != "none" and not
self.get_s3_endpoint_url():
+ # Keeps an unreachable STS endpoint from delaying the task by
more than a few seconds.
+ bounded = Config(connect_timeout=2, read_timeout=2,
retries={"total_max_attempts": 1})
+ config = self.aws_hook.config.merge(bounded)
+ sts_hook = StsHook(
+ aws_conn_id=self.aws_conn_id,
region_name=self.get_region_name(), config=config
Review Comment:
With `credential_chain` (the default), DuckDB ignores the keys and
`role_arn` of `aws_conn_id` and finds its own credentials, but `StsHook` uses
them. So the logged identity can be a different one from the one that got the
403, and people will look at the wrong IAM policy. Maybe use `aws_conn_id=None`
for `credential_chain`, so boto3 follows its default chain like DuckDB does?
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/duckdb.py:
##########
@@ -196,6 +200,25 @@ def get_s3_endpoint_url(self) -> str | None:
return explicit
return self.connection_extra.get("endpoint_url") or
self.aws_hook.conn_config.endpoint_url
+ def get_conn(self) -> DuckDBPyConnection:
+ """Return a DuckDB connection, and log the AWS identity once to help
debug S3 ``HTTP 403`` errors."""
+ conn = super().get_conn()
+ if self._aws_identity_looked_up:
+ return conn
+ self._aws_identity_looked_up = True
+ try:
+ if self.credential_strategy != "none" and not
self.get_s3_endpoint_url():
Review Comment:
`s3_endpoint_url` is also used for S3 VPC endpoints (see the docstring), and
those are real AWS. Should we only skip for non-AWS endpoints?
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/duckdb.py:
##########
@@ -196,6 +200,25 @@ def get_s3_endpoint_url(self) -> str | None:
return explicit
return self.connection_extra.get("endpoint_url") or
self.aws_hook.conn_config.endpoint_url
+ def get_conn(self) -> DuckDBPyConnection:
Review Comment:
nit: `configure_secrets` already computes `strategy` and `endpoint_url` and
returns early on `none`. Doing the lookup at the end of it would avoid checking
the same things twice.
##########
providers/amazon/tests/unit/amazon/aws/hooks/test_duckdb.py:
##########
@@ -305,3 +307,20 @@ def
test_explicit_argument_wins_over_the_connection_extra(self, aws_base_hook):
with mock.patch(GET_CONNECTION,
return_value=self.duckdb_connection(secret_name="from_extra")):
statement = secret_statement(AwsDuckDBHook(secret_name="explicit"))
assert statement.startswith("CREATE OR REPLACE SECRET explicit (")
+
+
+class TestAwsDuckDBHookIdentityLookup:
Review Comment:
Only the failure path is tested. Can you add tests for the happy path (ARN
logged) and the two skips (`credential_strategy="none"` and an endpoint
override)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
