workflow/dnet/dag_utils.py

@@ -0,0 +1,16 @@
+from airflow.hooks.base import BaseHook
+from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+
+def get_bucket_name(context: dict, hook: S3Hook, param_name: str):
+    bucket_name = context["params"][param_name]
+    if not bucket_name:
+        bucket_name = hook.extra_args['bucket_name']
+    return bucket_name
+
+
+def get_default_bucket():
+    hook = S3Hook("s3_conn", transfer_config_args={'use_threads': False})
+    try:
+        return hook.service_config['bucket_name']
+    except KeyError:
+        return ''