Michael McGuinness
168d197349
Add static type check * artifacts * dockerfileforall * examplemypystep * reqs * baseMain * oneperservice * somebasicchanges * deps * removedeps * removeerrs * singlecommand * lessargs * 312base * versionanderr * pythonversion * ithinkfinalerrs * numpyv * versionsandconditional * changesets * tryagain * lint * numpyagain * init * clearstreamlit * forcestreamlit Approved-by: Umang Mistry
388 lines
12 KiB
Python
388 lines
12 KiB
Python
import base64
|
|
import json
|
|
import snowflake.connector
|
|
import boto3
|
|
import logging
|
|
import os
|
|
|
|
|
|
"""
|
|
# Sample input events
|
|
|
|
doc_input_event = {
|
|
"operation": "insert",
|
|
"table": "DOCUMENT_LOGS",
|
|
"data": {
|
|
"DOCUMENT_ID" : "abc_xyz"
|
|
"BATCH_ID": "Batch_101",
|
|
"JOB_ID": "J123456",
|
|
"STAGE": "Processing",
|
|
"TEXTRACT_STATUS": "Success",
|
|
"BUCKET_NAME": "doc-bucket",
|
|
"FILE_NAME": "file1.pdf",
|
|
"FILE_PATH": "/documents/2023/",
|
|
"DOCUMENT_TYPE": "Report",
|
|
"PAYER_SIGNED": "True",
|
|
"PROVIDER_SIGNED": "False",
|
|
"GROUP_ID": "G100",
|
|
"CREATED_TIME": "2024-03-21 10:00:00",
|
|
"MODIFIED_TIME": "2024-03-21 10:00:00",
|
|
"CREATED_BY": "admin",
|
|
"MODIFIED_BY": "admin",
|
|
"ORIGINAL_FILE_EXTENSION": "pdf",
|
|
"NO_OF_PAGES": 10,
|
|
"FILE_SIZE": 1048576
|
|
}
|
|
}
|
|
|
|
|
|
|
|
doc_update_event = {
|
|
"operation": "update",
|
|
"table": "DOCUMENT_LOGS",
|
|
"data": {
|
|
"DOCUMENT_ID": "abc_xyz",
|
|
"TEXTRACT_STATUS": "Failed",
|
|
"MODIFIED_TIME": "2024-03-22 15:00:00",
|
|
"MODIFIED_BY": "admin"
|
|
}
|
|
}
|
|
|
|
|
|
|
|
batch_insert_event = {
|
|
"operation": "insert",
|
|
"table": "BATCH_LOGS",
|
|
"data": {
|
|
"CLIENT_ID": "C200",
|
|
"EXECUTION_START_TIME": "2024-03-21 09:00:00",
|
|
"NO_OF_DOCUMENTS": 150,
|
|
"USER_NAME": "batch_processor"
|
|
}
|
|
}
|
|
|
|
batch_update_event = {
|
|
"operation": "update",
|
|
"table": "BATCH_LOGS",
|
|
"data": {
|
|
"BATCH_ID": "batch_102",
|
|
"NO_OF_DOCUMENTS": 155,
|
|
"USER_NAME": "updated_processor"
|
|
}
|
|
}
|
|
|
|
client_insert_event = {
|
|
"operation": "insert",
|
|
"table": "CLIENT_LOGS",
|
|
"data": {
|
|
"CLIENT_ID": "CL300",
|
|
"CLIENT_NAME": "Acme Corporation",
|
|
"BUCKET_NAME": "acme-docs"
|
|
}
|
|
}
|
|
|
|
client_update_event = {
|
|
"operation": "update",
|
|
"table": "CLIENT_LOGS",
|
|
"data": {
|
|
"CLIENT_ID": "CL300",
|
|
"BUCKET_NAME": "new-acme-docs"
|
|
}
|
|
}
|
|
Operation can be: 'insert' or 'update'
|
|
Data is a dictionary with the columns and values to be inserted or updated
|
|
|
|
"""
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s", force=True)
|
|
logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
|
|
logging.getLogger("botocore").setLevel(logging.WARNING)
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
def get_secret(secrets_name: str):
|
|
"""Get credentials from Secret Manager as dict"""
|
|
secrets_manager = boto3.client("secretsmanager")
|
|
get_secret_value_response = secrets_manager.get_secret_value(SecretId=secrets_name)
|
|
|
|
if "SecretString" in get_secret_value_response:
|
|
secret_json = get_secret_value_response["SecretString"]
|
|
else:
|
|
secret_json = base64.b64decode(get_secret_value_response["SecretBinary"])
|
|
|
|
return json.loads(secret_json)
|
|
|
|
|
|
def get_snowflake_db_connection(secrets_name: str):
|
|
"""Create connection to Snowflake db."""
|
|
try:
|
|
con_params = get_secret(secrets_name)
|
|
account = con_params["account_locator"]
|
|
user = con_params["user"]
|
|
password = con_params["password"]
|
|
database = con_params["database"].upper()
|
|
warehouse = con_params["warehouse"]
|
|
role = con_params["role"]
|
|
logger.info(
|
|
f"Using credentials: account={account}, user={user}, password=***, database={database}, "
|
|
f"warehouse={warehouse}"
|
|
)
|
|
snowflake_connection = snowflake.connector.connect(
|
|
account=account,
|
|
user=user,
|
|
password=password,
|
|
database=database,
|
|
warehouse=warehouse,
|
|
role=role,
|
|
autocommit=True,
|
|
schema="STG",
|
|
)
|
|
logger.info(snowflake_connection)
|
|
cursor = snowflake_connection.cursor()
|
|
cursor.execute("SELECT * FROM STG.DIM_AUDIT LIMIT 10;")
|
|
logger.info(f"RESULT:: {cursor.fetchone()[0]}")
|
|
return snowflake_connection
|
|
except Exception as e:
|
|
return e
|
|
|
|
|
|
# Leaving this statement outside the lambda_handler function to reuse the connection
|
|
# Secret has been setup to use the logging service account
|
|
SECRET_MANAGER_NAME = os.environ.get("SECRET_MANAGER_NAME", "")
|
|
conn = get_snowflake_db_connection(SECRET_MANAGER_NAME)
|
|
# cur = conn.cursor()
|
|
# logger.info(f"CURSOR OBJECT: {cur}")
|
|
|
|
|
|
def construct_doc_insert_sql(data):
|
|
"""
|
|
Constructs the SQL for an insert operation
|
|
|
|
Sample return value:
|
|
INSERT INTO STG.DOCUMENT_LOGS (DOCUMENT_ID,BATCH_ID, JOB_ID, STAGE, TEXTRACT_STATUS, BUCKET_NAME, FILE_NAME, FILE_PATH, DOCUMENT_TYPE, PAYER_SIGNED, PROVIDER_SIGNED, GROUP_ID, CREATED_TIME, MODIFIED_TIME, CREATED_BY, MODIFIED_BY, ORIGINAL_FILE_EXTENSION, NO_OF_PAGES, FILE_SIZE)
|
|
VALUES ('doc_1212','batch_101', 'J123456', 'Processing', 'Success', 'doc-bucket', 'file1.pdf', '/documents/2023/', 'Report', True, False, 'G100', '2024-03-21 10:00:00', '2024-03-21 10:00:00', 'admin', 'admin', 'pdf', 10, 1048576);
|
|
"""
|
|
columns = ", ".join(data.keys())
|
|
values = ", ".join(
|
|
[
|
|
(
|
|
"'" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else str(value)
|
|
)
|
|
for value in data.values()
|
|
]
|
|
)
|
|
sql = f"INSERT INTO STG.DOCUMENT_LOGS ({columns}) VALUES ({values});"
|
|
return sql
|
|
|
|
|
|
def construct_doczy_pipeline_insert_sql(data):
|
|
"""
|
|
Constructs the SQL for an insert operation
|
|
|
|
Sample return value:
|
|
INSERT INTO STG.DOCUMENT_LOGS (DOCUMENT_ID,BATCH_ID, JOB_ID, STAGE, TEXTRACT_STATUS, BUCKET_NAME, FILE_NAME, FILE_PATH, DOCUMENT_TYPE, PAYER_SIGNED, PROVIDER_SIGNED, GROUP_ID, CREATED_TIME, MODIFIED_TIME, CREATED_BY, MODIFIED_BY, ORIGINAL_FILE_EXTENSION, NO_OF_PAGES, FILE_SIZE)
|
|
VALUES ('doc_1212','batch_101', 'J123456', 'Processing', 'Success', 'doc-bucket', 'file1.pdf', '/documents/2023/', 'Report', True, False, 'G100', '2024-03-21 10:00:00', '2024-03-21 10:00:00', 'admin', 'admin', 'pdf', 10, 1048576);
|
|
"""
|
|
columns = ", ".join(data.keys())
|
|
values = ", ".join(
|
|
[
|
|
(
|
|
"'" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else str(value)
|
|
)
|
|
for value in data.values()
|
|
]
|
|
)
|
|
sql = f"INSERT INTO STG.DOCZY_PIPELINE_RAW_OUTPUT ({columns}) VALUES ({values});"
|
|
return sql
|
|
|
|
|
|
def construct_doczy_pipeline_update_sql(data, document_id):
|
|
"""
|
|
Constructs the SQL for an update operation
|
|
|
|
Sample return value:
|
|
UPDATE STG.DOCUMENT_LOGS SET TEXTRACT_STATUS = 'Failed', MODIFIED_TIME = '2024-03-22 15:00:00', MODIFIED_BY = 'admin' WHERE DOCUMENT_ID = 1001;
|
|
"""
|
|
set_clauses = ", ".join(
|
|
[
|
|
(
|
|
f"{key} = '" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else f"{key} = {value}"
|
|
)
|
|
for key, value in data.items()
|
|
]
|
|
)
|
|
sql = f"UPDATE STG.DOCZY_PIPELINE_RAW_OUTPUT SET {set_clauses} WHERE DOCUMENT_ID = '{document_id}';"
|
|
logger.info(f"Executing the following logging SQL statement: {sql}")
|
|
return sql
|
|
|
|
|
|
def construct_doc_update_sql(data, document_id):
|
|
"""
|
|
Constructs the SQL for an update operation
|
|
|
|
Sample return value:
|
|
UPDATE STG.DOCUMENT_LOGS SET TEXTRACT_STATUS = 'Failed', MODIFIED_TIME = '2024-03-22 15:00:00', MODIFIED_BY = 'admin' WHERE DOCUMENT_ID = 1001;
|
|
"""
|
|
set_clauses = ", ".join(
|
|
[
|
|
(
|
|
f"{key} = '" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else f"{key} = {value}"
|
|
)
|
|
for key, value in data.items()
|
|
]
|
|
)
|
|
sql = f"UPDATE STG.DOCUMENT_LOGS SET {set_clauses} WHERE DOCUMENT_ID = '{document_id}';"
|
|
logger.info(f"Executing the following logging SQL statement: {sql}")
|
|
return sql
|
|
|
|
|
|
def construct_batch_insert_sql(data):
|
|
"""
|
|
Constructs the SQL for an insert operation
|
|
|
|
Sample return value:
|
|
INSERT INTO STG.BATCH_LOGS (BATCH_ID,CLIENT_ID, EXECUTION_START_TIME, NO_OF_DOCUMENTS, USER_NAME) VALUES ('C200', '2024-03-21 09:00:00', 150, 'batch_processor');
|
|
"""
|
|
columns = ", ".join(data.keys())
|
|
values = ", ".join(
|
|
[
|
|
(
|
|
"'" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else str(value)
|
|
)
|
|
for value in data.values()
|
|
]
|
|
)
|
|
sql = f"INSERT INTO STG.BATCH_LOGS ({columns}) VALUES ({values});"
|
|
return sql
|
|
|
|
|
|
def construct_client_insert_sql(data):
|
|
"""
|
|
Constructs the SQL for an insert operation
|
|
|
|
Sample return value:
|
|
INSERT INTO STG.CLIENT_LOGS (CLIENT_ID, CLIENT_NAME, BUCKET_NAME) VALUES ('CL300', 'Acme Corporation', 'acme-docs');
|
|
"""
|
|
columns = ", ".join(data.keys())
|
|
values = ", ".join(
|
|
[
|
|
(
|
|
"'" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else str(value)
|
|
)
|
|
for value in data.values()
|
|
]
|
|
)
|
|
sql = f"INSERT INTO STG.CLIENT_LOGS ({columns}) VALUES ({values});"
|
|
return sql
|
|
|
|
|
|
def construct_client_update_sql(data, client_id):
|
|
"""
|
|
Constructs the SQL for an update operation
|
|
|
|
Sample return value:
|
|
UPDATE STG.CLIENT_LOGS SET BUCKET_NAME = 'new-acme-docs' WHERE CLIENT_ID = 'CL300';
|
|
"""
|
|
set_clauses = ", ".join(
|
|
[
|
|
(
|
|
f"{key} = '" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else f"{key} = {value}"
|
|
)
|
|
for key, value in data.items()
|
|
]
|
|
)
|
|
sql = f"UPDATE STG.CLIENT_LOGS SET {set_clauses} WHERE CLIENT_ID = '{client_id}';"
|
|
return sql
|
|
|
|
|
|
def construct_batch_update_sql(data, batch_id):
|
|
"""
|
|
Constructs the SQL for an update operation
|
|
|
|
Sample return value:
|
|
UPDATE STG.BATCH_LOGS SET NO_OF_DOCUMENTS = 155, USER_NAME = 'updated_processor' WHERE BATCH_ID = 1;
|
|
"""
|
|
set_clauses = ", ".join(
|
|
[
|
|
(
|
|
f"{key} = '" + str(value).replace("'", "''") + "'"
|
|
if isinstance(value, str)
|
|
else f"{key} = {value}"
|
|
)
|
|
for key, value in data.items()
|
|
]
|
|
)
|
|
sql = f"UPDATE STG.BATCH_LOGS SET {set_clauses} WHERE BATCH_ID = '{batch_id}';"
|
|
return sql
|
|
|
|
|
|
# Main Lambda handler
|
|
def lambda_handler(event, context):
|
|
# Extract operation type and payload from event
|
|
# res = cur.execute("SELECT CURRENT_DATABASE();").fetchone()
|
|
# logger.info(f"current database: {res}")
|
|
operation = event["operation"] # 'insert' or 'update'
|
|
data = event["data"]
|
|
table = event["table"]
|
|
|
|
with conn.cursor() as cur:
|
|
try:
|
|
if table == "DOCUMENT_LOGS":
|
|
if operation == "insert":
|
|
sql = construct_doc_insert_sql(data)
|
|
elif operation == "update":
|
|
document_id = data.pop("DOCUMENT_ID", None)
|
|
sql = construct_doc_update_sql(data, document_id)
|
|
|
|
elif table == "BATCH_LOGS":
|
|
if operation == "insert":
|
|
sql = construct_batch_insert_sql(data)
|
|
elif operation == "update":
|
|
batch_id = data.pop("BATCH_ID", None)
|
|
sql = construct_batch_update_sql(data, batch_id)
|
|
|
|
elif table == "DOCZY_PIPELINE_RAW_OUTPUT":
|
|
if operation == "insert":
|
|
sql = construct_doczy_pipeline_insert_sql(data)
|
|
# elif operation == 'update':
|
|
# batch_id = data.pop('BATCH_ID', None)
|
|
# sql = construct_doczy_pipeline_update_sql(data, batch_id)
|
|
|
|
elif table == "CLIENT_LOGS":
|
|
if operation == "insert":
|
|
sql = construct_client_insert_sql(data)
|
|
elif operation == "update":
|
|
client_id = data.pop("CLIENT_ID", None)
|
|
sql = construct_client_update_sql(data, client_id)
|
|
else:
|
|
raise ValueError("Unsupported table.")
|
|
|
|
response = cur.execute(sql)
|
|
logger.info(f"Executing the following logging SQL statement: {sql}")
|
|
return {"statusCode": 200, "body": json.dumps("Operation successful")}
|
|
|
|
except Exception as e:
|
|
return {"statusCode": 400, "body": json.dumps(str(e))}
|
|
finally:
|
|
# conn.close()
|
|
# Need to close the connection to avoid reaching the limit of open connections
|
|
# However, we need the conn to be in hot state for subsequent concurrent executions
|
|
# Need to decide the best approach to handle this
|
|
pass
|