def _setup_source_conn(self, source_conn_id, source_bucket_name=None):
"""
Retrieve connection based on source_conn_id. In case of s3 it also configures the bucket.
Validates that connection id belongs to supported connection type.
:param source_conn_id:
:param source_bucket_name:
"""
self.source_conn = BaseHook.get_hook(source_conn_id)
self.source_conn_id = source_conn_id
# Workaround for getting hook in case of s3 connection
# This is needed because get_hook silently returns None for s3 connections
# See https://issues.apache.org/jira/browse/AIRFLOW-2316 for more info
connection = BaseHook._get_connection_from_env(source_conn_id)
self.log.info(connection.extra_dejson)
if connection.conn_type == 's3':
self.log.info("Setting up s3 connection {0}".format(source_conn_id))
self.source_conn = S3Hook(aws_conn_id=source_conn_id)
# End Workaround
if source_bucket_name is None:
raise AttributeError("Missing source bucket for s3 connection")
self.source_bucket_name = source_bucket_name
if not isinstance(self.source_conn, DbApiHook) and not isinstance(self.source_conn, S3Hook):
raise AttributeError(
"Only s3_csv, local and sql connection types are allowed, not {0}".format(type(self.source_conn)))
def _setup_dest_conn(self, dest_conn_id, results_bucket_name, results_dest_name):
"""
Setup results connection. Retrieves s3 connection and makes sure we've got location details (bucket, filename)
:param dest_conn_id:
:param results_bucket_name:
:param results_dest_name:
"""
conn = BaseHook._get_connection_from_env(dest_conn_id)
if conn.conn_type != 's3':
raise AttributeError(
"Only s3 is allowed as a results destination, not {0}".format(conn.conn_type))
self.dest_conn = S3Hook(aws_conn_id=dest_conn_id)
self.dest_conn_id = dest_conn_id
if results_bucket_name is None or results_dest_name is None:
raise AttributeError("Specify bucket name and key name to store results")
self.results_bucket_name = results_bucket_name
self.results_dest_name = results_dest_name
