from pyspark.sql import SparkSession,DataFrame
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
# from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from pyspark.sql.functions import expr, col , base64
import boto3
import numpy as np
from awsglue.dynamicframe import DynamicFrame
from pyspark.sql import functions as F
from pyspark.sql.types import DateType, StringType
from pyspark.sql.types import StructType,StructField,StringType,IntegerType, BooleanType
from datetime import datetime
import yaml
import regex as re
import json
import hashlib
import random
from typing import Dict, List, Tuple,Union
spark = SparkSession.builder.appName('ops-mask-test').getOrCreate()
sc = spark.sparkContext
glueContext = GlueContext(sc)
job = Job(glueContext)
# create df from catalog
def create_df_from_catalog(catalog_table_name, latest_year, latest_month, latest_day):
df = glueContext.create_data_frame_from_catalog(database=source_catalog_database,
table_name=catalog_table_name,
transformation_ctx='dyf',
push_down_predicate=f"(year=={latest_year} and month=={latest_month} and day=={latest_day})"
)
return df
# get all partition
def get_all_partitions(catalog_database, catalog_table_name):
partitions = glue_client.get_partitions(
DatabaseName=catalog_database,
TableName=catalog_table_name
)
partitions_values = partitions['Partitions']
partitions_sorted_all = sorted((list(map(lambda d: d['Values'], partitions_values))))
return partitions_sorted_all
# Remaining days to sync
def remaining_dates_to_sync(source_catalog_table):
partition_source_all = get_all_partitions(source_catalog_database, source_catalog_table) # returns all partitions in bronze
try:
partition_destination_all = get_all_partitions(destination_catalog_database, source_catalog_table)
except:
return partition_source_all
partition_destination_latest = max(partition_destination_all)
remaining_dates = [date for date in partition_source_all if date > partition_destination_latest]
if not remaining_dates:
return [date for date in partition_source_all if date >= partition_destination_latest]
return remaining_dates
# Convert to dynamic df and write to s3 as parquet ### This will auto update the data catalog and create table if not present
def write_data(df, catalog_table_name):
lastest_source_partition = df.select('year', 'month', 'day').collect()[0].asDict()
year = lastest_source_partition["year"]
month = lastest_source_partition["month"]
day = lastest_source_partition["day"]
dynamic_df = DynamicFrame.fromDF(df.coalesce(1), glueContext, catalog_table_name)
glueContext.purge_s3_path(destination_s3_path+'/'+catalog_table_name +'/year='+year+'/month='+month+'/day='+day, {"retentionPeriod": 0})
sink = glueContext.getSink(
connection_type="s3",
updateBehavior="UPDATE_IN_DATABASE",
path=destination_s3_path+'/'+catalog_table_name,
enableUpdateCatalog=True,
partitionKeys=["year", "month", "day"])
sink.setFormat("glueparquet")
sink.setCatalogInfo(catalogDatabase=destination_catalog_database,catalogTableName=catalog_table_name)
sink.writeFrame(dynamic_df)
def get_anonymize_dict(DatabaseName: str, TableName: str) -> Dict:
response = glue_client.get_table(DatabaseName=DatabaseName, Name=TableName)
anonymize_dict = {
'hash': [],
'randomize': []
}
for column in response['Table']['StorageDescriptor']['Columns']:
for keys, values in column.items():
# if keys == 'Parameters' and len(values)>0:
if keys == 'Parameters' and bool(values):
new_dict = column['Parameters']
dict_value = new_dict.get('Class', None)
if dict_value == 'PII':
anonymize_dict['hash'].append(column['Name'])
elif dict_value == 'FIN':
fin_tuple=(column['Name'],int(new_dict['Lower']),int(new_dict['Upper']))
anonymize_dict['randomize'].append(fin_tuple)
return anonymize_dict
def anonymize(df: DataFrame, schema = Dict[str, Union[List[str], List[Tuple[str, int, int]]]])->DataFrame:
"""
Anonymizes sensitive data in a DataFrame based on a provided schema.
Args:
df (DataFrame): The DataFrame containing the sensitive data.
schema (Dict[str, List[str]]): A schema specifying the columns to anonymize and the anonymization methods to apply.
Returns:
DataFrame: A new DataFrame with the specified columns anonymized according to the schema.
Raises:
None.
Example:
schema = {
'mask': ['email', 'phone_number'],
'randomize': ['age']
}
"""
for keys, values in schema.items():
if keys == 'hash':
for value in values:
# print(value)
df = df.withColumn(value, F.md5(value))
elif keys=='randomize':
for value in values:
column=value[0]
lower=value[1]
upper= value[2]
# print(column,lower,upper)
# print("randomize: " +value)
randomize_udf = F.udf(lambda: np.random.randint(lower, upper), IntegerType())
df = df.withColumn(column, randomize_udf())
return df
############### FOR JSON START START START #######################
def get_anonymize_json_dict(DatabaseName: str, TableName: str) -> Dict:
response = glue_client.get_table(DatabaseName=DatabaseName, Name=TableName)
anonymize_json_dict = {
'hash': [],
'randomize': []
}
for column in response['Table']['StorageDescriptor']['Columns']:
for keys, values in column.items():
# if keys == 'Parameters' and len(values)>0:
if keys == 'Parameters' and bool(values):
new_dict = column['Parameters']
dict_value = new_dict.get('Type', None)
if dict_value == 'JSON':
anonymize_json_dict['hash'].extend(new_dict['PII'].split(','))
anonymize_json_dict['randomize'] = [tuple(x.strip().replace('(', '').replace(')', '').split(',')) for x in new_dict['FIN'].split('),')]
anonymize_json_dict['randomize'] = [(item[0], int(item[1]), int(item[2])) for item in anonymize_json_dict['randomize']]
anonymize_json_dict['json_col']=column['Name']
return anonymize_json_dict
def extract_keys_to_anonymize (dict_provided):
keys_to_anonymize = []
for key, value in dict_provided.items():
if key == 'hash':
keys_to_anonymize.extend(value)
elif key == 'randomize':
keys_to_anonymize.extend([x[0] for x in value])
return keys_to_anonymize
def anonymize_json_value(value, key, keys_to_anonymize, anonymize_json_dict):
if value is None:
return None
if key in keys_to_anonymize:
if key in anonymize_json_dict['hash']:
return hashlib.md5(str(value).encode('utf-8')).hexdigest()
elif key in [x[0] for x in anonymize_json_dict['randomize']]:
for x in anonymize_json_dict['randomize']:
if key == x[0]:
return random.randint(int(x[1]), int(x[2]))
elif isinstance(value, list):
return [anonymize_json_value(item, key, keys_to_anonymize, anonymize_json_dict) for item in value]
elif isinstance(value, dict):
return {k: anonymize_json_value(val, k, keys_to_anonymize, anonymize_json_dict) for k, val in value.items()}
return value
def anonymize_json(json_data, keys_to_anonymize, anonymize_json_dict):
if isinstance(json_data, str):
json_data = json.loads(json_data)
return json.dumps(anonymize_json_value(json_data, None, keys_to_anonymize, anonymize_json_dict))
############### FOR JSON END END END #######################
def rename_columns(df):
for col in df.columns:
new_col = re.sub(r'\s+|[^A-Za-z0-9]', '_', col)
df = df.withColumnRenamed(col, new_col)
return df
def get_tables_list(database_name):
table_list_response = glue_client.get_tables(
DatabaseName=source_catalog_database)['TableList']
return sorted([table_dict['Name'] for table_dict in table_list_response], reverse=False)
def handle_writing_task(source_catalog_database):
tables_list = get_tables_list(source_catalog_database)
for catalog_table_name in tables_list:
print(f'Processing table : {catalog_table_name}')
remaining_dates = remaining_dates_to_sync(catalog_table_name)
print(f'Remaining Dates: {remaining_dates}')
for remaining_date_calculated in remaining_dates:
latest_year = remaining_date_calculated[0]
latest_month = remaining_date_calculated[1]
latest_day = remaining_date_calculated[2]
print(f'Processing {catalog_table_name} for {remaining_date_calculated}')
df = rename_columns(create_df_from_catalog(
catalog_table_name, latest_year, latest_month, latest_day))
try:
anonymize_dict = get_anonymize_dict(
source_catalog_database, catalog_table_name)
anonymize_json_dict=get_anonymize_json_dict(
source_catalog_database, catalog_table_name)
keys_to_anonymize=extract_keys_to_anonymize(anonymize_json_dict)
json_col=anonymize_json_dict['json_col']
anonymize_json_udf = F.udf(lambda x: anonymize_json(x, keys_to_anonymize, anonymize_json_dict), StringType())
print(keys_to_anonymize)
print(anonymize_dict)
print(anonymize_json_dict)
df_ano = anonymize(df, anonymize_dict)
df_ano=df_ano.withColumn(json_col, anonymize_json_udf(json_col))
except:
df_ano = df
df_ano.show(2)
# df_ano.select(json_col).show(2,truncate=False)
#### Uncomment the line below to write anomymize data
write_data(df_ano, catalog_table_name)
print(f'Completed {catalog_table_name} for {remaining_date_calculated}')