import logging
import datetime
from pymorri.util.bigquery_client import BigQueryClient
from pymorri.util.serverless_utils import get_or_exception
from pymorri.util.aws_utils import pass_event_via_s3, chunked_list, get_event_from_sns
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# If we pass event via SNS only `get_event_from_sns' decorator will be used and `pass_event_via_s3` + `chunked_list`
# should exit early without changes. If we pass event via StepFunctions (S3) `pass_event_via_s3` + `chunked_list`
# decorators will be used and 'get_event_from_sns' should exit early without changes
[docs]@get_event_from_sns
@pass_event_via_s3
@chunked_list("bq_update_data")
def lambda_handler(event, context):
"""
This Lambda will be called in one of the 2 ways. It will be either triggered by SNS when message is placed
in the sns-euw1-${env}-digital-storepick-events topic and the message attribute attached to the message is
eventType=logRowToBq or called by an upstream Lambda which is part of the step function.
This Lambda is for updating bq_records in BigQuery.
Structured event to craft SQL UPDATE query contains:
- config_file - config file related to the service
- bq_update_data - details required to build bigquery sql query (dataset_and_table is optional - if not provided then same will be used from Config)
This function works both for UPDATE and INSERT:
- if `condition` == INSERT, then it will INSERT data based on `columns_to_update` dictionary
- if `condition` != INSERT, then it will UPDATE data based on `columns_to_update` dictionary AND specified condition (for WHERE)
Args:
event (dict) - required:
"config_file" - config file related to the service
"bq_update_data" - data for update
Example:
.. code-block:: python
event = {
"config_file": "storepick-nhs-discount.json",
"bq_update_data" : [
{
"dataset_and_table" : "store_pick_discount.nhs_discount",
"condition" : "status=\"PENDING\"",
"columns_to_update" : {
"status" : "FAILED",
"message" : "test message",
},
},
]
}
The following query will be crafted based on the above example:
UPDATE `store_pick_discount.nhs_discount` SET status="FAILED", message="test message" WHERE retailerCustomerId="aaa-bbb"
Returns:
event (dict): Updated input event with 'bq_records' list containing rows read from BigQuery.
Example:
.. code-block:: python
resp = {
"config_file": "storepick-nhs-discount.json",
"bq_get_data_query": "SQL query used to update data in BigQuery",
"bq_records": [
{
'key1': 'value1',
'key2': 'value2',
...
},
{
'key1': 'value3',
'key2': 'value4',
...
},
]
}
Raises:
MoriException: MoriException containing details about exception
GoogleCloudBased exceptions: Any error message received from GCP API
"""
config_file = get_or_exception(event, "config_file")
bq_update_data = get_or_exception(event, "bq_update_data")
bq_client = BigQueryClient(
configuration=config_file,
)
columns_to_update = []
rows_to_insert = []
condition = ""
for update_row in bq_update_data:
dataset_and_table = "__TABLE_FROM_CONFIG__"
get_dataset_table = update_row.get("dataset_and_table")
if get_dataset_table:
dataset_and_table = f"`{get_dataset_table}`"
condition = get_or_exception(update_row, "condition")
columns_to_update = get_or_exception(update_row, "columns_to_update")
if condition == "INSERT":
rows_to_insert.append(columns_to_update)
else:
formatted_update_data = ""
for column, value in columns_to_update.items():
if isinstance(value, str):
formatted_update_data += f'{column}="{value}", '
else:
formatted_update_data += f'{column}={value}, '
# remove trailing ',' char
formatted_update_data = formatted_update_data.rstrip(", ")
query = f"UPDATE {dataset_and_table} SET {formatted_update_data} WHERE {condition}"
logger.info(f"query to UPDATE: '{query}'")
# get statuses and react
_ = bq_client.query(query)
if rows_to_insert:
logger.info(f"INSERTing rows to BigQuery => nr of rows: '{len(rows_to_insert)}'")
# using 'write' function as it's performance is much better than separate queries
bq_client.write(rows_to_insert)
del event["bq_update_data"]
event['nr_of_bq_updates'] = len(bq_update_data)
return event
if __name__ == "__main__":
import sys
condition = "INSERT"
if len(sys.argv) > 1 and sys.argv[1] == "INSERT":
condition = f'(status=\"SUCCESS\" or status=\"FAILED\") AND promotionName="{sys.argv[1]}"'
event = {
"config_file": "storepick-nhs-discount.json",
"bq_update_data": [
{
"dataset_and_table": "store_pick_discount.nhs_discount",
"condition": condition,
"columns_to_update": {
"retailerCustomerId": "98765432",
"clubId": 100,
"status": "PENDING",
"message": "",
"timestamp": datetime.datetime.now().isoformat(),
"operation": "addCustomerToGroup",
},
},
]
}
event.update({"skip_pass_event_via_s3": True}) # skip s3 for single lambda invokations only
lambda_handler(event, {})
event = {
"Records": [
{
"EventVersion": "1.0",
"EventSubscriptionArn": "arn:aws:sns:us-east-2:123456789012:sns-lambda:21be56ed-a058-49f5-8c98"
"-aedd2564c486",
"EventSource": "aws:sns",
"Sns": {
"SignatureVersion": "1",
"Timestamp": "2019-01-02T12:45:07.000Z",
"Signature": "tcc6faL2yUC6dgZdmrwh1Y4cGa/ebXEkAi6RibDsvpi+tE/1+82j...65r==",
"SigningCertUrl": "https://sns.us-east-2.amazonaws.com/SimpleNotificationService"
"-ac565b8b1a6c5d002d285f9598aa1d9b.pem",
"MessageId": "95df01b4-ee98-5cb9-9903-4c221d41eb5e",
"Message": '{"config_file": "storepick-nhs-discount.json", "bq_update_data": [{ '
'"dataset_and_table": "store_pick_discount.nhs_discount", "condition": "('
'status=\\"SUCCESS\\" OR status=\\"FAILED\\")", "columns_to_update": { "status": '
'"PENDING", "message": ""}}] }',
"MessageAttributes": {
"eventType": {
"Type": "String",
"Value": "logRowToBq"
}
},
"Type": "Notification",
"UnsubscribeUrl": "https://sns.us-east-2.amazonaws.com/?Action=Unsubscribe&SubscriptionArn"
"=arn:aws:sns:us-east-2:123456789012:test-lambda:21be56ed-a058-49f5-8c98"
"-aedd2564c486",
"TopicArn": "arn:aws:sns:us-east-2:123456789012:sns-lambda",
"Subject": "TestInvoke"
}
}
]
}
lambda_handler(event, {})