You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ds...@apache.org on 2023/01/04 09:27:10 UTC
[airflow] branch main updated: Don't call get_connection from provide_bucket_name (#28716)
This is an automated email from the ASF dual-hosted git repository.
dstandish pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9ab9c18d68 Don't call get_connection from provide_bucket_name (#28716)
9ab9c18d68 is described below
commit 9ab9c18d687b61a8b86bde2767369697288d8082
Author: Daniel Standish <15...@users.noreply.github.com>
AuthorDate: Wed Jan 4 01:27:01 2023 -0800
Don't call get_connection from provide_bucket_name (#28716)
We can used the cached connection instead of retrieving the conn object (maybe causing roundtrip to metastore) every time.
---
airflow/providers/amazon/aws/hooks/s3.py | 6 ++----
airflow/providers/amazon/aws/utils/connection_wrapper.py | 2 ++
tests/providers/amazon/aws/hooks/test_s3.py | 10 +++-------
tests/providers/amazon/aws/utils/test_connection_wrapper.py | 2 +-
4 files changed, 8 insertions(+), 12 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
index 1e2235f24a..dc8b0f2f69 100644
--- a/airflow/providers/amazon/aws/hooks/s3.py
+++ b/airflow/providers/amazon/aws/hooks/s3.py
@@ -57,10 +57,8 @@ def provide_bucket_name(func: T) -> T:
if "bucket_name" not in bound_args.arguments:
self = args[0]
- if self.aws_conn_id:
- connection = self.get_connection(self.aws_conn_id)
- if connection.schema:
- bound_args.arguments["bucket_name"] = connection.schema
+ if self.conn_config and self.conn_config.schema:
+ bound_args.arguments["bucket_name"] = self.conn_config.schema
return func(*bound_args.args, **bound_args.kwargs)
diff --git a/airflow/providers/amazon/aws/utils/connection_wrapper.py b/airflow/providers/amazon/aws/utils/connection_wrapper.py
index f6ae3fdf01..237109fa75 100644
--- a/airflow/providers/amazon/aws/utils/connection_wrapper.py
+++ b/airflow/providers/amazon/aws/utils/connection_wrapper.py
@@ -103,6 +103,7 @@ class AwsConnectionWrapper(LoggingMixin):
conn_type: str | None = field(init=False, default=None)
login: str | None = field(init=False, repr=False, default=None)
password: str | None = field(init=False, repr=False, default=None)
+ schema: str | None = field(init=False, repr=False, default=None)
extra_config: dict[str, Any] = field(init=False, repr=False, default_factory=dict)
# AWS Credentials from connection.
@@ -156,6 +157,7 @@ class AwsConnectionWrapper(LoggingMixin):
self.conn_type = conn.conn_type or "aws"
self.login = conn.login
self.password = conn.password
+ self.schema = conn.schema or None
self.extra_config = deepcopy(conn.extra_dejson)
if self.conn_type.lower() == "s3":
diff --git a/tests/providers/amazon/aws/hooks/test_s3.py b/tests/providers/amazon/aws/hooks/test_s3.py
index ff231482da..8b57247e63 100644
--- a/tests/providers/amazon/aws/hooks/test_s3.py
+++ b/tests/providers/amazon/aws/hooks/test_s3.py
@@ -472,13 +472,9 @@ class TestAwsS3Hook:
def test_function(self, bucket_name=None):
return bucket_name
- fake_s3_hook = FakeS3Hook()
-
- test_bucket_name = fake_s3_hook.test_function()
- assert test_bucket_name == mock_get_connection.return_value.schema
-
- test_bucket_name = fake_s3_hook.test_function(bucket_name="bucket")
- assert test_bucket_name == "bucket"
+ hook = FakeS3Hook()
+ assert hook.test_function() == "test_bucket"
+ assert hook.test_function(bucket_name="bucket") == "bucket"
def test_delete_objects_key_does_not_exist(self, s3_bucket):
# The behaviour of delete changed in recent version of s3 mock libraries.
diff --git a/tests/providers/amazon/aws/utils/test_connection_wrapper.py b/tests/providers/amazon/aws/utils/test_connection_wrapper.py
index 853245bfa0..d1238e1ffa 100644
--- a/tests/providers/amazon/aws/utils/test_connection_wrapper.py
+++ b/tests/providers/amazon/aws/utils/test_connection_wrapper.py
@@ -89,10 +89,10 @@ class TestAwsConnectionWrapper:
assert wrap_conn.extra_config is not mock_conn.extra_dejson
# `extra_config` is a same object that return by `extra_dejson`
assert wrap_conn.extra_config is wrap_conn.extra_dejson
+ assert wrap_conn.schema == "mock-schema"
# Check that not assigned other attributes from airflow.models.Connection to wrapper
assert not hasattr(wrap_conn, "host")
- assert not hasattr(wrap_conn, "schema")
assert not hasattr(wrap_conn, "port")
# Check that Wrapper is True if assign connection