jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Added business config objs & client config

Browse Source
This commit is contained in:
Umang Mistry
2024-03-27 16:06:06 -05:00
parent d1d83f7842
commit cc59a1ccc1
4 changed files with 254 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lambda/textract_logs_lambda.py
+123 -130
View File
@@ -1,54 +1,142 @@
import json
import snowflake.connector
import boto3
import snowflake.connector
import logging
"""
Sample input event for Insert operation:
{
# Sample input events
doc_input_event = {
"operation": "insert",
"table": "DOCUMENT_LOGS",
"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",
"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": 204800
"FILE_SIZE": 1048576
}
}
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
}
}
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):
"""
@@ -133,10 +221,8 @@ def lambda_handler(event, context):
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)
@@ -159,108 +245,15 @@ def lambda_handler(event, context):
sql = construct_client_update_sql(data, client_id)
else:
raise ValueError("Unsupported table.")
print(sql)
# cursor.execute(sql)
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
# 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)