rework replace_illegal_characters function

Author: timburke-hackit

scripts/jobs/data_and_insight/icaseworks_ingest_to_raw.py

@@ -13,9 +13,16 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-arg_keys = ['region_name', 's3_endpoint', 's3_target_location', 's3_staging_location', 'target_database',
-            'target_table', 'secret_name']
-partition_keys = ['import_year', 'import_month', 'import_day', 'import_date']
+arg_keys = [
+    "region_name",
+    "s3_endpoint",
+    "s3_target_location",
+    "s3_staging_location",
+    "target_database",
+    "target_table",
+    "secret_name",
+]
+partition_keys = ["import_year", "import_month", "import_day", "import_date"]
 
 import boto3
 
@@ -41,14 +48,20 @@
 
 ### Functions ###
 
+
 def remove_illegal_characters(string):
-    """Removes illegal characters from string"""
-    regex_list = [['=', ""], ['\/', "_"], ['+', "-"]]
-    for r in regex_list:
-        clean_string = re.sub(string=string,
-                              pattern="[{}]".format(r[0]),
-                              repl=r[1])
-    return clean_string
+    """Removes illegal characters from string by replacing:
+    = with empty string
+    / with _
+    + with -
+    """
+    replacements = {"=": "", "/": "_", "+": "-"}
+
+    result = string
+    for old, new in replacements.items():
+        result = result.replace(old, new)
+
+    return result
 
 
 def encode_json(json_string):
@@ -62,32 +75,34 @@ def encode_json(json_string):
 def create_signature(header, payload, secret):
     """Encode JSON string"""
     # hashed header, hashed payload, string secret
-    unsigned_token = header + '.' + payload
+    unsigned_token = header + "." + payload
     # secret_access_key = base64.b64decode(unsigned_token) #TODO is this used anywhere??
-    key_bytes = bytes(secret, 'utf-8')
-    string_to_sign_bytes = bytes(unsigned_token, 'utf-8')
-    signature_hash = hmac.new(key_bytes, string_to_sign_bytes, digestmod=hashlib.sha256).digest()
+    key_bytes = bytes(secret, "utf-8")
+    string_to_sign_bytes = bytes(unsigned_token, "utf-8")
+    signature_hash = hmac.new(
+        key_bytes, string_to_sign_bytes, digestmod=hashlib.sha256
+    ).digest()
     encoded_signature = base64.b64encode(signature_hash)
-    encoded_signature = encoded_signature.decode('utf-8')
+    encoded_signature = encoded_signature.decode("utf-8")
     encoded_signature = remove_illegal_characters(encoded_signature)
     return encoded_signature
 
 
 def get_token(url, encoded_header, encoded_payload, signature, headers):
     """Get token"""
     assertion = encoded_header + "." + encoded_payload + "." + signature
-    data = f'assertion={assertion}&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer'
-    print(f'Data : {data}')
+    data = f"assertion={assertion}&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer"
+    print(f"Data : {data}")
     response = requests.request("POST", url, headers=headers, data=data)
     return response
 
 
 def get_icaseworks_report_from(report_id, fromdate, auth_headers, auth_payload):
     report_url = "https://hackneyreports.icasework.com/getreport?"
-    request_url = f'{report_url}ReportId={report_id}&Format=json&From={fromdate}'
-    print(f'Request url: {request_url}')
+    request_url = f"{report_url}ReportId={report_id}&Format=json&From={fromdate}"
+    print(f"Request url: {request_url}")
     r = requests.request("GET", request_url, headers=auth_headers, data=auth_payload)
-    print(f'Status Code: {r.status_code}')
+    print(f"Status Code: {r.status_code}")
     return r
 
 
@@ -96,25 +111,25 @@ def get_report_fromtime(report_id, timestamp_to_call, auth_headers, auth_payload
 
     now = str(datetime.now())
     timetouse = now[:19]
-    print(f'TimeNow: {timetouse}')
+    print(f"TimeNow: {timetouse}")
 
-    request_url = f'{report_url}ReportId={report_id}&Format=json&FromTime={timestamp_to_call}&UntilTime={timetouse}'
-    print(f'Request url: {request_url}')
+    request_url = f"{report_url}ReportId={report_id}&Format=json&FromTime={timestamp_to_call}&UntilTime={timetouse}"
+    print(f"Request url: {request_url}")
     r = requests.request("GET", request_url, headers=auth_headers, data=auth_payload)
-    print(f'Status Code: {r.status_code}')
+    print(f"Status Code: {r.status_code}")
     return r
 
 
 def dump_dataframe(response, location, filename):
-    df = pd.DataFrame.from_dict(response.json(), orient='columns')
+    df = pd.DataFrame.from_dict(response.json(), orient="columns")
 
-    df['import_year'] = datetime.today().year
-    df['import_month'] = datetime.today().month
-    df['import_day'] = datetime.today().day
-    df['import_date'] = datetime.today().strftime('%Y%m%d')
+    df["import_year"] = datetime.today().year
+    df["import_month"] = datetime.today().month
+    df["import_day"] = datetime.today().day
+    df["import_date"] = datetime.today().strftime("%Y%m%d")
 
-    print(f'Database: {target_database}')
-    print(f'Table: {target_table}')
+    print(f"Database: {target_database}")
+    print(f"Table: {target_table}")
 
     # write to s3
     wr.s3.to_parquet(
@@ -124,28 +139,27 @@ def dump_dataframe(response, location, filename):
         database=target_database,
         table=target_table,
         mode="overwrite_partitions",
-        partition_cols=partition_keys
+        partition_cols=partition_keys,
     )
-    print(f'Dumped Dataframe {df.shape} to {s3_target_location}')
-    logger.info(f'Dumped Dataframe {df.shape} to {s3_target_location}')
+    print(f"Dumped Dataframe {df.shape} to {s3_target_location}")
+    logger.info(f"Dumped Dataframe {df.shape} to {s3_target_location}")
 
 
 def get_latest_timestamp(table_dict):
     #   try:
-    print(f'Getting max timestamp')
+    print(f"Getting max timestamp")
     # 2025-01-05T15:06:16
 
     sql_query = 'select max("casestatustouchtimes_lastupdatedtime") as latest from "data-and-insight-raw-zone"."icaseworks_foi"'
 
-    conn = connect(s3_staging_dir=s3_staging_location,
-                   region_name=region_name)
+    conn = connect(s3_staging_dir=s3_staging_location, region_name=region_name)
 
     df = pd.read_sql_query(sql_query, conn)
     latest_date = df.iloc[0, 0]
     latest_date = latest_date.replace("T", " ")
 
-    print(f'dataframe outputting')
-    print(f'Time Found: {latest_date}')
+    print(f"dataframe outputting")
+    print(f"Time Found: {latest_date}")
 
     return latest_date
 
@@ -155,87 +169,95 @@ def get_latest_timestamp(table_dict):
 #     print(f'No Data Found. Will use {date_to_return}')
 #     return date_to_return
 
+
 def authenticate_icaseworks(api_key, secret):
     url = "https://hackney.icasework.com/token"
 
-    headers = {
-        'Content-Type': 'application/x-www-form-urlencoded'
-    }
+    headers = {"Content-Type": "application/x-www-form-urlencoded"}
 
     header_object = {"alg": "HS256", "typ": "JWT"}
 
     # Create Header
     header_object = str(header_object).replace("'", '"').replace(" ", "")
     header = encode_json(header_object)
-    print(f'Header: {header}')
+    print(f"Header: {header}")
 
     # Create payload
     current_unix_time = int(time.time())
     str_time = str(current_unix_time)
-    payload_object = {
-        "iss": api_key,
-        "aud": url,
-        "iat": str_time
-    }
-    payload_object = str(payload_object).replace("'", '"').replace(" ",
-                                                                   "")  # can we do a dict-to-string function for this and the header
+    payload_object = {"iss": api_key, "aud": url, "iat": str_time}
+    payload_object = (
+        str(payload_object).replace("'", '"').replace(" ", "")
+    )  # can we do a dict-to-string function for this and the header
 
     payload = encode_json(str(payload_object))
-    print(f'Created Payload: {payload}')
+    print(f"Created Payload: {payload}")
 
     # Create Signature
     signature = create_signature(header, payload, secret)
-    print(f'Created Signature: {signature}')
+    print(f"Created Signature: {signature}")
 
     # Get assertion
     assertion = header + "." + payload + "." + signature
-    print(f'assertion: {assertion}')
+    print(f"assertion: {assertion}")
 
     # Get response
-    response = get_token(url=url, encoded_header=header, encoded_payload=payload, signature=signature, headers=headers)
+    response = get_token(
+        url=url,
+        encoded_header=header,
+        encoded_payload=payload,
+        signature=signature,
+        headers=headers,
+    )
     print(response)
 
     # Get token
-    auth_token = response.json().get('access_token')
-    print(f'auth token: {auth_token}')
+    auth_token = response.json().get("access_token")
+    print(f"auth token: {auth_token}")
 
     # Create auth header for API Calls and auth payload
 
-    authorization = f'Bearer {auth_token}'
+    authorization = f"Bearer {auth_token}"
     print(authorization)
 
     auth_payload = []
 
     # Note: I don't know how to generate the below cookie. That is extracted using postman. Not sure how to recreate this at all
-    auth_headers = {
-        'Authorization': authorization
-    }
-    print(f'')
+    auth_headers = {"Authorization": authorization}
+    print(f"")
     return auth_payload, auth_headers
 
 
 def get_data(table_dict, date_to_call, auth_headers, auth_payload):
     dict_to_call = table_dict
 
-    print(f'Pulling report for {dict_to_call["name"]}')
+    print(f"Pulling report for {dict_to_call['name']}")
     case_id_report_id = dict_to_call["reportid"]
-    case_id_list = get_report_fromtime(case_id_report_id, date_to_call, auth_headers, auth_payload)
+    case_id_list = get_report_fromtime(
+        case_id_report_id, date_to_call, auth_headers, auth_payload
+    )
     # print(f'Type of case_id_list {type(case_id_list)}')
 
-    dict_to_call["DF"] = case_id_list  # This will append the response to the DF column in the dictionary
+    dict_to_call["DF"] = (
+        case_id_list  # This will append the response to the DF column in the dictionary
+    )
 
     dump_dataframe(dict_to_call["DF"], dict_to_call["location"], date_to_call)
 
 
 def get_data_from(table_dict, date_to_call, auth_headers, auth_payload):
     dict_to_call = table_dict
 
-    print(f'Pulling report for {dict_to_call["name"]}')
+    print(f"Pulling report for {dict_to_call['name']}")
     case_id_report_id = dict_to_call["reportid"]
-    case_id_list = get_icaseworks_report_from(case_id_report_id, date_to_call, auth_headers, auth_payload)
+    case_id_list = get_icaseworks_report_from(
+        case_id_report_id, date_to_call, auth_headers, auth_payload
+    )
     # print(f'Type of case_id_list {type(case_id_list)}')
 
-    dict_to_call["DF"] = case_id_list  # This will append the response to the DF column in the dictionary
+    dict_to_call["DF"] = (
+        case_id_list  # This will append the response to the DF column in the dictionary
+    )
 
     dump_dataframe(dict_to_call["DF"], dict_to_call["location"], date.today())
 
@@ -249,35 +271,41 @@ def retrieve_credentials_from_secrets_manager(secrets_manager_client, secret_nam
 
 ### main function ##
 
+
 def main():
-    secrets_manager_client = boto3.client('secretsmanager')
-    api_credentials_response = retrieve_credentials_from_secrets_manager(secrets_manager_client, secret_name)
-    api_credentials = json.loads(api_credentials_response['SecretString'])
+    secrets_manager_client = boto3.client("secretsmanager")
+    api_credentials_response = retrieve_credentials_from_secrets_manager(
+        secrets_manager_client, secret_name
+    )
+    api_credentials = json.loads(api_credentials_response["SecretString"])
     api_key = api_credentials.get("api_key")
-    print(f'Api_key: {api_key}')
+    print(f"Api_key: {api_key}")
     secret = api_credentials.get("secret")
-    print(f'Secret: {secret}')
+    print(f"Secret: {secret}")
     auth_payload, auth_headers = authenticate_icaseworks(api_key, secret)
 
     list_of_datadictionaries = [
         #   {"name":"Corrective Actions", "reportid":188769, "full_ingestion":False, "location":"/content/drive/MyDrive/iCaseworks/Corrective_Actions/"},
         #   {"name":"Classifications", "reportid":188041, "full_ingestion":True, "location":"/content/drive/MyDrive/iCaseworks/classifications/"},
-        {"name": "FOI Requests", "reportid": 199549, "full_ingestion": False,
-         "location": "s3://dataplatform-stg-raw-zone/data-and-insight/icaseworks_foi/"}
+        {
+            "name": "FOI Requests",
+            "reportid": 199549,
+            "full_ingestion": False,
+            "location": "s3://dataplatform-stg-raw-zone/data-and-insight/icaseworks_foi/",
+        }
     ]
 
     for data_dict in list_of_datadictionaries:
         location = data_dict["location"]
 
-        if data_dict['full_ingestion'] == False:
-
+        if data_dict["full_ingestion"] == False:
             date_to_track_from = get_latest_timestamp(data_dict)
-            print(f'Starting calls from {date_to_track_from}')
+            print(f"Starting calls from {date_to_track_from}")
 
             get_data(data_dict, date_to_track_from, auth_headers, auth_payload)
         else:
             get_data_from(data_dict, "2020-09-01", auth_headers, auth_payload)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()