import json import snowflake.connector import boto3 import snowflake.connector import logging """ # 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" } } 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__) 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, autocommit=True) logger.info(snowflake_connection) 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 conn = get_snowflake_db_connection('doczy-dev-db-svc-acc') cur = conn.cursor() 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'] 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 == '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.") logger.info(f"Executing the following logging SQL statement: {sql}") cur.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() # 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