17 lines
513 B
Python
17 lines
513 B
Python
|
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 ''
|