Files
doczyai-pipelines/lambda/textract_logs_lambda.py
T

267 lines
7.9 KiB
Python
Raw Normal View History

import json
import snowflake.connector
"""
Sample input event for Insert operation:
{
"operation": "insert",
"data": {
"BATCH_ID": 12345,
"JOB_ID": "job_67890",
"STAGE": "Preprocessing",
"TEXTRACT_STATUS": "Pending",
"BUCKET_NAME": "my-bucket",
"FILE_NAME": "document.pdf",
"FILE_PATH": "/path/to/document.pdf",
"DOCUMENT_TYPE": "Type A",
"PAYER_SIGNED": False,
"PROVIDER_SIGNED": True,
"GROUP_ID": "group_123",
"CREATED_TIME": "2023-01-01 12:00:00",
"MODIFIED_TIME": "2023-01-01 12:00:00",
"CREATED_BY": "user_1",
"MODIFIED_BY": "user_1",
"ORIGINAL_FILE_EXTENSION": "pdf",
"NO_OF_PAGES": 10,
"FILE_SIZE": 204800
}
}
Sample input event for Update operation:
{
"operation": "update",
"data": {
"DOCUMENT_ID": 1001,
"TEXTRACT_STATUS": "Completed",
"MODIFIED_TIME": "2023-01-02 13:00:00",
"MODIFIED_BY": "user_2",
"NO_OF_PAGES": 12,
"FILE_SIZE": 304800
}
}
Operation can be: 'insert' or 'update'
Data is a dictionary with the columns and values to be inserted or updated
"""
def construct_doc_insert_sql(data):
"""
Constructs the SQL for an insert operation
Sample return value:
INSERT INTO STG.DOCUMENT_LOGS (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 (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_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};"
return sql
def construct_batch_insert_sql(data):
"""
Constructs the SQL for an insert operation
Sample return value:
INSERT INTO STG.BATCH_LOGS (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
operation = event['operation'] # 'insert' or 'update'
data = event['data']
table = event['table']
# Get conn from Secrets Manager
try:
# with conn.cursor() as cursor:
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 == '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.")
print(sql)
# cursor.execute(sql)
return {'statusCode': 200, 'body': json.dumps('Operation successful')}
except Exception as e:
return {'statusCode': 400, 'body': json.dumps(str(e))}
finally:
# conn.close()
pass
# Sample input events
doc_input_event = {
"operation": "insert",
"table": "DOCUMENT_LOGS",
"data": {
"BATCH_ID": 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": 1001,
"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": 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"
}
}
lambda_handler(doc_input_event, None)
lambda_handler(doc_update_event, None)
lambda_handler(batch_insert_event, None)
lambda_handler(batch_update_event, None)
lambda_handler(client_insert_event, None)
lambda_handler(client_update_event, None)