Multi-Region Disaster Recovery for QLDB

"""
This Lambda is responsible for exporting the Ledger from the QLDB database to S3
"""
import os
from datetime import datetime, timedelta
import boto3

dr_bucket_name = os.getenv("DR_BUCKET_NAME")
ledger_name = os.getenv("LEDGER_NAME")
region = os.getenv("AWS_REGION")
role_arn = os.getenv("ROLE_ARN")
qldb_export_trigger_hour = os.getenv("QLDB_EXPORT_TRIGGER_HOUR")

# Configure the AWS SDK with your credentials and desired region
session = boto3.Session(region_name=region)
qldb_client = session.client("qldb")

def handler(event, context) -> tuple[int, dict]:
    """
    Lambda entry point


    event: object
    Event passed when the lambda is triggered


    context: object
    Lambda Context


    return: tuple[int, dict]
    Success code and an empty object
    """

    time_end = datetime.now().replace(minute=0, second=0, microsecond=0)
    time_start = time_end - timedelta(hours=int(qldb_export_trigger_hour))
    # Specify the QLDB ledger name and export configuration
    export_config = {
        "Bucket": dr_bucket_name,
        "Prefix": "qldb-dr/",
        "EncryptionConfiguration": {"ObjectEncryptionType": "SSE_S3"},
    }

    # Trigger the export by calling the ExportJournalToS3 API
    response = qldb_client.export_journal_to_s3(
        Name=ledger_name,
        S3ExportConfiguration=export_config,
        InclusiveStartTime=time_start,
        ExclusiveEndTime=time_end,
        RoleArn=role_arn,
        OutputFormat="JSON",
    )

    # Print the export job ID
    export_job_id = response["ExportId"]
    print("Export Job ID: " + export_job_id)

    return 200, {}

{
   ...
   "transactionId":"A2mztA7Su6k5dRMQhOHN5C",
   "transactionInfo":{
      "statements":[
         {
            "statement":"SELECT * FROM information_schema.user_tables",
            "startTime":"2023-05-24T17:00:33.988Z"
         }
      ]
   }
},

{
   ...
   "transactionId":"A2mztA8BDHN5R62i10guzh",
   "transactionInfo":{
      "statements":[
         {
            "statement":"INSERT INTO test_account2 ?",
            "startTime":"2023-05-24T17:11:17.098Z"
         }
      ],
      ...
   "revisions":[
      {
         "data":{
            "uuid":"9eff5e3f-eb90-11ed-9a6e-0a3efd619f29",
            "accountNo":268,
            "accountName":"Miscellaneous Expenses."
         },
 ...
}

"""
This Lambda is responsible for splitting up the load from
the QLDB exports distributes the load to be processed
by multiple Lambdas instances
"""
import os
import boto3


# Configure the S3 and SQS clients
s3_client = boto3.client("s3")
sqs_client = boto3.client("sqs")

# Retrieve the name of the source S3 bucket and target SQS queue from the event
source_bucket = os.getenv("DR_BUCKET_NAME")
target_queue_url = os.getenv("SQS_QUEUE_URL")


def handler(event, context) -> tuple[int, dict]:
    """
    Lambda entry point


    event: object
    Event passed when the lambda is triggered


    context: object
    Lambda Context


    return: tuple[int, dict]
    Success code and an empty object
    """

    # List all objects in the source bucket
    try:
        object_keys = []
        paginator = s3_client.get_paginator("list_objects_v2")
        pages = paginator.paginate(Bucket=source_bucket)

        for page in pages:
            for obj in page["Contents"]:
                object_keys.append(obj["Key"])
    except Exception as e:
        print(f"DR Export Error when listing objects in bucket: {source_bucket}
{e}")
        raise

    # Push each object key to the SQS queue
    for object_key in object_keys:
        if object_key.endswith(".json"):
            try:
                sqs_client.send_message(
                    QueueUrl=target_queue_url,
                    MessageBody=f"s3://{source_bucket}/{object_key}",
                )
            except Exception as e:
                print(f"DR Export Error when publishing message to SQS queue {e}")
            raise

    return 200, {}

"""
This Lambda is responsible for restoring the Ledger from the S3 backup bucket to QLDB database
"""
import os
import json

import awswrangler as wr
import botocore
from pyqldb.config.retry_config import RetryConfig
from pyqldb.driver.qldb_driver import QldbDriver

region = os.getenv("AWS_REGION")
ledger_name = os.getenv("LEDGER_NAME")


def handler(event, context) -> tuple[int, dict]:
    """
    Lambda entry point


    event: object
    Event passed when the lambda is triggered


    context: object
    Lambda Context


    return: tuple[int, dict]
    Success code and an empty object
    """

    # Reading from SQS queue
    for record in event["Records"]:
        s3_path = record["body"]

        file_dataframe = wr.s3.read_json(path=s3_path, lines=True)
        df_to_json_array = json.loads(file_dataframe.to_json(orient="records"))

        docs_to_insert = []
        for item in df_to_json_array:
            for statement in item["transactionInfo"]["statements"]:
                if (
                    statement["statement"]
                    .upper()
                    .startswith(("INSERT INTO", "UPDATE", "DELETE"))
                ):
                    docs_to_insert.append(
                        {
                            "statement": statement["statement"],
                            "data": item["revisions"][0]["data"],
                        }
                    )

            driver = QldbDriver(
                ledger_name=ledger_name,
                retry_config=RetryConfig(retry_limit=0),
                config=botocore.config.Config(
                    retries={"max_attempts": 0},
                    read_timeout=10,
                    connect_timeout=10,
                ),
                region_name=region,
            )

        try:
            for doc in docs_to_insert:
                driver.execute_lambda(
                    lambda x: x.execute_statement(doc["statement"], doc["data"])
                )
        except Exception as e:
            print("DR Restore Error in file:" + s3_path + str(e))
            raise

        print("Restored file:" + s3_path)

    return 200, {}