def _execute_batches(self, batches: List):
results = []
LOG.info("Executing cassandra batches")
for batch in batches:
results.append(self._execute_batch(batch))
LOG.info("finished %s batches", len(results))
return results
def _show_result(self, execution_id, max_result_size=1000):
results = self._get_query_result(execution_id, max_result_size)
column_info = results['ResultSet']['ResultSetMetadata']['ColumnInfo']
headers = [h['Name'].encode('utf-8') for h in column_info]
LOG.info(headers)
csv_writer = csv.writer(sys.stdout, quoting=csv.QUOTE_ALL)
csv_writer.writerows([[val['VarCharValue'] for val in row['Data']]
for row in results['ResultSet']['Rows']])
def create_table(self, table_settings):
"""
Create a table from given settings
Args:
table_settings (dict): Dictionary of settings to create table
Returns: None
"""
table_sql = self._build_create_table_sql(table_settings)
LOG.info(table_sql)
self.run_query(table_sql)
def _state_manager_connect(self):
LOG.info("Connecting to Cassandra")
conn = CassandraConnectionManager(
self.__settings.etl_state_manager_connection)
conn.setup_connection(self.__settings.etl_state_manager_keyspace)
LOG.info("Cassandra connection established")
sync_etl_state_table()
def get_all_keys(self, key_prefix: str) -> List[str]:
"""
Sense all keys under a given key prefix
Args:
key_prefix (str): the key prefix under which all files will be sensed
Returns: List[str]
"""
LOG.info("sensing files from s3://%s/%s ", self.bucket, key_prefix)
metadata = self.get_object_metadata(key_prefix)
lines = [file.key for file in metadata]
LOG.info("found %s s3 keys", len(lines))
return lines
def _upsert_data_frame(self, data_frame):
if self.__settings.destination_batch_size > 1:
LOG.info("Going to upsert batches of size %s", self.__settings.destination_batch_size)
result = self._get_cassandra_util().upsert_dataframe_in_batches(
dataframe=data_frame,
table=self.__settings.destination_table,
batch_size=self.__settings.destination_batch_size)
else:
LOG.info("Going to upsert one row at a time")
result = self._get_cassandra_util().upsert_dataframe(
dataframe=data_frame,
table=self.__settings.destination_table)
return result
def delete_recursive(self, key_prefix: str) -> None:
"""
Recursively delete all keys with given prefix from the named bucket
Args:
key_prefix (str): Key prefix under which all files will be deleted
Returns: NA
"""
if not key_prefix.endswith("/"):
key_prefix = f"{key_prefix}/"
LOG.info("Recursively deleting s3://%s/%s", self.bucket, key_prefix)
response = self.get_resource().Bucket(self.bucket).objects.filter(
Prefix=key_prefix).delete()
LOG.info(response)
def rename_file(self, key: str, new_file_name: str) -> None:
"""
Rename a file on s3
Args:
key: Current key of the file
new_file_name: target file name
Returns: None
"""
s3 = self.get_resource()
full_new_file_path = key.rpartition('/')[0] + '/' + new_file_name
LOG.info("Renaming source: %s to %s", key, full_new_file_path)
s3.Object(self.bucket, full_new_file_path).copy_from(
CopySource={'Bucket': self.bucket, 'Key': key})
s3.Object(self.bucket, key).delete()
def select_dataframe(self, query) -> DataFrame:
"""
execute SQL query using Airflow mysql hook and retrieve data in pandas data frame
Args:
query (str): Mysql compliant query string
Returns: DataFrame
"""
LOG.info("Executing \n %s", query)
with closing(self._conn_mgr.get_conn()) as connection:
df = sql.read_sql(query, connection)
LOG.info("Sql Data frame size: %s", df.shape[0])
LOG.debug(df.head(2))
return df
def _multi_process_upload_file(settings: AwsConnectionSettings, filename: str, bucket: str,
key: str) -> None:
"""
A standalone copy of the method making it simple to pickle in a multi processing pool
Args:
settings: the s3 connection settings to use for upload
filename: local file name of the file to be uploaded.
bucket: the s3 bucket to upload file to.
key: the s3 key to use while uploading the file
Returns: None
"""
LOG.info("Uploading File %s to s3://%s/%s", filename, bucket, key)
S3Util(
conn=AwsConnectionManager(settings),
bucket=bucket
).upload_file(local_file_path=filename, key=key)
def get_page_as_list_of_dict(page: dict) -> List[OrderedDict]:
"""
Converts a list of entries from google adwords response into a list of Ordered Dictionaries
Args:
page (dict): the response page from google adwords api
Returns: List[dict]
"""
result = []
if 'entries' in page:
entries = page['entries']
# These entries are a list of zeep Objects that need conversion to Dict
result = [zeep_object_to_dict(entry) for entry in entries]
LOG.debug("The result from the adword API: %s", result)
else:
LOG.info('No entries were found.')
return result
def parse_and_export_msgs(self, list_of_msgs, interval):
"""
Converts messages to a pandas dataframe and then exports to s3
Args:
list_of_msgs (list(Kafka Message Object)): List of msg objects
interval (int): Rounding interval for the temporal partitioning
Returns: None
"""
good_data, bad_data = convert_msgs_to_dictionary(list_of_msgs)
self._partition_data_and_upload_to_s3(good_data, interval)
self._partition_data_and_upload_to_s3(bad_data, interval)
LOG.info("Data Upload Complete")
def execute(self, query: str, values: Optional[Tuple] = None) -> None:
"""
execute standalone SQL statements in a mysql database
Args:
query (str): Mysql compliant query string
values (Tuple): n tuple for substituting values in the query
Returns: None
"""
LOG.info("Executing \n %s", query)
conn = self._conn_mgr.get_conn()
cur = conn.cursor()
if values:
cur.execute(query, values)
else:
cur.execute(query)
conn.commit()
def delete_recursive_match_suffix(self, key_prefix: str, suffix: str) -> None:
"""
Recursively delete all keys with given key prefix and suffix from the bucket
Args:
key_prefix (str): Key prefix under which all files will be deleted.
suffix (str): suffix of the subset of files in the given prefix directory to be deleted
Returns: None
"""
if not key_prefix:
raise ValueError("key_prefix must not be empty")
if not suffix:
raise ValueError("suffix must not be empty")
s3 = self.get_resource()
for obj in s3.Bucket(self.bucket).objects.filter(Prefix=key_prefix):
if obj.key.endswith(suffix):
LOG.info("deleting s3://%s/%s", self.bucket, obj.key)
response = obj.delete()
LOG.info("Response: %s ", response)
def download_directory(self, source_key: str, file_suffix: str, local_directory: str) -> None:
"""
Download an entire directory from s3 onto local file system
Args:
source_key (str): key prefix of the directory to be downloaded from s3
file_suffix (str): suffix to filter a subset under the source_key to be downloaded
local_directory (str): local absolute path to store all the files
Returns: None
"""
s3 = self.get_resource()
LOG.info("Downloading s3://%s/%s to %s", self.bucket, source_key, local_directory)
for obj in s3.Bucket(self.bucket).objects.filter(Prefix=source_key):
key_path = obj.key.split("/")
if obj.key.endswith(file_suffix):
filename = f"{local_directory}/{key_path[-1]}"
self.download_file(
local_file_path=filename,
key=obj.key)
def watch_query(self, execution_id, poll_frequency=10):
"""
Watch the query execution for a given execution id in Athena
Args:
execution_id: the execution id of an Athena Auery
poll_frequency (int): Freq in seconds to poll for the query status using Athen API
Returns: dictionary of status from Athena
"""
LOG.info("Watching query with execution id - %s", execution_id)
while True:
athena = self.get_client()
stats = athena.get_query_execution(QueryExecutionId=execution_id)
status = stats['QueryExecution']['Status']['State']
if status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
LOG.info("Query Completed %s", stats)
return stats
time.sleep(poll_frequency)
def upload_all(self) -> (List[dict], List[dict], List[dict]):
"""
Upload all files from the Athena table onto AdWords offline conversion
Returns:
verification_issues List[dict]: a tuple of lists outlining any
verification failures
successes List[dict]: The responses for successful uploads to
the Google Adwords API
failures List[dict]: The responses for failed uploads to the
Google Adwords API
"""
verification_issues, successes, failures = [], [], []
for key in self.list_source_files():
issues, success, fail = \
self._process_data_frame(self.get_data_frame(key))
verification_issues.extend(issues)
successes.extend(success)
failures.extend(fail)
if len(verification_issues) > 0:
LOG.warning("There were %s verification failures",
len(verification_issues))
LOG.debug("All verification failures: \n %s", verification_issues)
if len(failures) > 0:
LOG.warning(
"There were %s failures uploading to the adwords "
"API", len(failures))
LOG.info("Sample Failure: \n %s", failures[0])
LOG.debug("All failures: \n %s", failures)
LOG.info(
"There were %s records successfully uploaded from a total of %s submitted items",
len(successes),
len(successes) + len(failures) + len(verification_issues))
return verification_issues, successes, failures
def prepare_batches(self,
prepared_statement: PreparedStatement,
tuples: List[tuple],
batch_size: int) -> List:
"""
Prepares a list of cassandra batched Statements out of a list of tuples and prepared
statement
Args:
prepared_statement (PreparedStatement): the statement to be used for batching.
tuples (list[tuple]): the data to be inserted.
batch_size (int): limit on the number of prepared statements in the batch.
Returns: list[BatchStatement]
"""
batches = []
LOG.debug("Preparing cassandra batches out of rows")
batches_of_tuples = _chunk_list(tuples, batch_size)
for tpl in batches_of_tuples:
batch = self._prepare_batch(prepared_statement, tpl)
batches.append(batch)
LOG.info("created %s batches out of list of %s tuples", len(batches), len(tuples))
return batches
def download_strings_from_directory(self, key_prefix: str, encoding: str = UTF8) -> List[str]:
"""
Read lines from s3 files
Args:
key_prefix: the key prefix under which all files will be read
encoding: the character encoding to use for encoding / decoding content
Returns: List[str]
"""
s3 = self.get_resource()
bucket = s3.Bucket(name=self.bucket)
lines = []
LOG.info("reading files from s3://%s/%s", self.bucket, key_prefix)
file_metadata = bucket.objects.filter(Prefix=key_prefix)
for file in file_metadata:
obj = s3.Object(self.bucket, file.key)
data = obj.get()['Body'].read().decode(encoding)
lines.append(data.splitlines())
# Flatten the list of lists
flat_lines = [item for sublist in lines for item in sublist]
LOG.info("read %s lines from %s s3 files", len(flat_lines), len(lines))
return flat_lines
def run_query(self, query_string, return_result=False):
"""
General purpose query executor that submits an athena query, then uses the execution id
to poll and monitor the
sucess of the query. and optionally return the result.
Args:
query_string (string): The string contianing valid athena query
return_result (boolean): Boolean flag to turn on results
Returns (boolean): if return_result = True then returns result dictionary, else None
"""
athena = self.get_client()
output_location = "s3://{bucket}/{key}".format(
bucket=self.output_bucket, key=self.output_key)
LOG.info(
"executing query \n%s \non database - %s with results location %s",
query_string, self.database, output_location)
response = athena.start_query_execution(
QueryString=query_string,
QueryExecutionContext={'Database': self.database},
ResultConfiguration={'OutputLocation': output_location})
execution_id = response['QueryExecutionId']
stats = self.watch_query(execution_id)
LOG.info("athena response %s", response)
if stats['QueryExecution']['Status']['State'] == 'SUCCEEDED':
LOG.info("Query execution id - %s SUCCEEDED", execution_id)
if return_result:
return self._get_query_result(execution_id)
else:
raise ValueError("Query exited with {} state because {}".format(
stats['QueryExecution']['Status']['State'],
stats['QueryExecution']['Status']['StateChangeReason']))
return None
def upload_directory(self,
source_directory: str,
extension: str,
target_key: str,
overwrite: bool = True,
rename: bool = True) -> None:
"""
Upload a local file directory to s3
Args:
source_directory (str): Local source directory's absolute path.
extension (str): the file extension of files in that directory to be uploaded.
target_key (str): Target location on the s3 bucket for files to be uploaded.
overwrite (bool): overwrites files on s3 if set to true.
rename (bool): renames the file when uploading to s3 if set to true.
Returns: None
"""
if overwrite:
LOG.info("Cleaning existing files on s3")
self.delete_recursive(f"{target_key}/")
LOG.info("searching for files to upload in %s", source_directory)
path_list = Path(source_directory).glob(f'**/*.{extension}')
itr = 0
upload_data = []
for path in path_list:
path_in_str = str(path)
filename = os.path.basename(path_in_str)
if rename:
filename = f"file-{str(uuid.uuid4())}.{extension}"
destination_key = f"{target_key}/{filename}"
itr = itr + 1
upload_data += [(self.conn.settings, path_in_str, self.bucket, destination_key)]
Pool().starmap(_multi_process_upload_file, upload_data) # Use all available cores
LOG.info("Upload of directory complete at s3://%s/%s", self.bucket, target_key)
def _build_create_table_sql(self, table_settings):
exists = _construct_table_exists_ddl(table_settings["exists"])
partitions = _construct_table_partition_ddl(
table_settings["partitions"])
table_properties = _construct_table_properties_ddl(
table_settings.get("skip_headers", False),
table_settings["storage_format_selector"].lower(),
table_settings["encryption"])
sql = """
CREATE EXTERNAL TABLE {exists} {table}(
{columns}
)
{partitions}
ROW FORMAT SERDE
'{row_format_serde}'
STORED AS INPUTFORMAT
'{inputformat}'
OUTPUTFORMAT
'{outputformat}'
LOCATION
's3://{s3_bucket}/{s3_dir}'
{table_properties}
""".format(
table=table_settings["table"],
exists=exists,
columns=zip_columns(table_settings["columns"]),
partitions=partitions,
row_format_serde=self.storage_format_lookup[
table_settings["storage_format_selector"]]["row_format_serde"],
inputformat=self.storage_format_lookup[
table_settings["storage_format_selector"]]["inputformat"],
outputformat=self.storage_format_lookup[
table_settings["storage_format_selector"]]["outputformat"],
s3_bucket=table_settings["s3_bucket"],
s3_dir=table_settings["s3_dir"],
table_properties=table_properties)
LOG.info("Query from table settings:\n %s", sql)
return sql
def create_athena_table(self) -> None:
"""
Creates an athena table on top of the transferred data
Returns: None
"""
athena_util = self._get_athena_util()
s3_util = self._get_s3_util()
keys = s3_util.get_keys(key_prefix=self._final_target_prefix)
LOG.debug("gathered files transferred under this ETL %s", keys)
if keys:
data = s3_util.download_parquet_as_dataframe(keys[0])
LOG.info(
"Downloaded parquet file from s3 to construct Athena create table statement: %s "
"\n made dataframe of shape %s", keys[0], data.shape)
if self.__settings.target_table_ddl_progress:
athena_util.drop_table(self.__settings.target_table)
athena_table_settings = self._construct_athena_table_settings(data)
athena_util.create_table(table_settings=athena_table_settings)
else:
raise ValueError(
"No Data has been uploaded to target directory, please load data first, "
"before creating Athena table")
def get_keys_modified_in_range(self,
key_prefix: str,
start_date: arrow,
end_date: arrow) -> List[str]:
"""
Sense if there were any files changed or added in the given time period under the given key
prefix and return a list of keys
Args:
key_prefix: the key prefix under which all files will be sensed
start_date: start of the duration in which the s3 objects were modified
end_date: end of the duration in which the s3 objects were modified
Returns: List[str]
"""
LOG.info("sensing files from s3://%s/%s \n between %s to %s", self.bucket, key_prefix,
start_date, end_date)
metadata = self.get_object_metadata(key_prefix)
lines = []
for file in metadata:
if start_date < arrow.get(file.last_modified) <= end_date:
lines += [file.key]
LOG.info("found %s s3 files changed", len(lines))
return lines
def move_recursive(self,
source_dir: str,
destination_dir: str,
delete_after_copy: bool = True) -> None:
"""
recursively move files on s3 to a new location on the same bucket
Args:
source_dir: Source key prefix representing the directory to move
destination_dir: destination key prefix
delete_after_copy: removes the files from source after successful copy if set to true
Returns: None
"""
s3 = self.get_resource()
bucket = s3.Bucket(self.bucket)
for obj in bucket.objects.filter(Prefix=source_dir):
# replace the prefix
new_key = destination_dir + obj.key[len(source_dir):]
LOG.info("Moving s3 object from : \n%s \nto: \n%s", obj.key, new_key)
new_obj = bucket.Object(new_key)
new_obj.copy({'Bucket': self.bucket, 'Key': obj.key})
if delete_after_copy:
self.delete_recursive(source_dir)
def instantiate_producer(self):
"""
Try to connect to the Kafka bootstrap server. We include
functionality to allow failure to connect to the queue to happen
silently
Returns (Producer): Kafka Producer
"""
try:
LOG.debug("Instantiating Producer")
self.producer = create_producer(self.producer_config)
# Check the connection works by polling for messages
LOG.debug("Polling Queue")
self.producer.poll(3)
LOG.info("Succesfully polled Kafka Queue")
except Exception as exception:
self.producer = None
LOG.error("Kafka Producer failed to instantiate: \n %s", exception)
if self.raise_exception_on_failed_connection:
raise NoProducerInstantiatedError()