feat: Add Lambda function to trigger initial SQL scripts (#2)

* task: New lambda function

* task: Git ignore

* task: Git ignore fix

* task: Added Lambda function to trigger Postgres setup

* fix: Unnecessary file

.gitignore

@@ -62,3 +62,7 @@ npm-debug.log*
 yarn-debug.log*
 yarn-error.log*
 pnpm-debug.log*
+
+# Temp Lambda files:
+backend/lambda_*
+lambda_*/

__main__.py

@@ -14,6 +14,9 @@
 import os
 import pulumi
 import json
+import subprocess
+import shutil
+import sys
 import pulumi_aws as aws
 import pulumi_docker as docker
 from dotenv import load_dotenv
@@ -340,7 +343,6 @@
         "cidr_blocks": ["0.0.0.0/0"]
     }]
 )
-
 langserve_database_cluster = aws.rds.Cluster(
     "langserve-database",
     engine=aws.rds.EngineType.AURORA_POSTGRESQL,
@@ -517,6 +519,89 @@
         "Name": f"{pulumi_project}-{pulumi_stack}",
     })
 
+# Define the IAM role for the Lambda function
+sql_migration_lambda_role = aws.iam.Role("langserve-lambda-role",
+    assume_role_policy="""{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Action": "sts:AssumeRole",
+            "Principal": {
+                "Service": "lambda.amazonaws.com"
+            },
+            "Effect": "Allow",
+            "Sid": ""
+        }
+    ]
+    }""")
+
+sql_migration_lambda_policy = aws.iam.RolePolicy(
+    "langserve-lambda-policy",
+    role=sql_migration_lambda_role.id,
+    policy="""{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Effect": "Allow",
+            "Action": [
+                "rds:*",
+                "logs:*",
+                "cloudwatch:*"
+            ],
+            "Resource": "*"
+        }
+    ]
+    }""")
+
+
+aws.iam.RolePolicyAttachment(
+    "lambda-vpc-access-execution-role-attachment",
+    role=sql_migration_lambda_role.name,
+    policy_arn="arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole"
+)
+
+FUNCTION_CODE_DIR = "./backend/sql"
+LAMBDA_PACKAGE_DIR = "./lambda_package"
+
+# Clean up any existing package directory
+if os.path.exists(LAMBDA_PACKAGE_DIR):
+    shutil.rmtree(LAMBDA_PACKAGE_DIR)
+
+os.makedirs(LAMBDA_PACKAGE_DIR, exist_ok=True)
+
+# Install psycopg2 into the package directory
+subprocess.check_call(
+    [sys.executable, "-m", "pip", "install", "psycopg2-binary", "-t", LAMBDA_PACKAGE_DIR, "--platform", "manylinux2014_x86_64", "--no-deps", "--only-binary", ":all:"]
+)
+
+# Copy the function code to the package directory
+shutil.copytree(FUNCTION_CODE_DIR, os.path.join(LAMBDA_PACKAGE_DIR, "code"))
+
+# Package the Lambda function and its dependencies
+shutil.make_archive(f"{LAMBDA_PACKAGE_DIR}/lambda_package", 'zip', LAMBDA_PACKAGE_DIR)
+
+sql_migration_lambda_function = aws.lambda_.Function(
+    'sql-migrations-lambda-function',
+    code=pulumi.FileArchive(f"{LAMBDA_PACKAGE_DIR}/lambda_package.zip"),
+    runtime='python3.9',
+    role=sql_migration_lambda_role.arn,
+    handler='code.sql_migration_lambda_function.lambda_handler',
+    environment=aws.lambda_.FunctionEnvironmentArgs(
+        variables={
+            'POSTGRES_HOST': langserve_cluster_instance.endpoint,
+            'POSTGRES_PORT': langserve_cluster_instance.port,
+            'POSTGRES_DB': langserve_database_cluster.database_name,
+            'POSTGRES_USER': langserve_database_cluster.master_username,
+            'POSTGRES_PASSWORD': langserve_database_cluster.master_password
+        }
+    ),
+    vpc_config=aws.lambda_.FunctionVpcConfigArgs(
+        subnet_ids=[langserve_subnet1.id, langserve_subnet2.id],
+        security_group_ids=[database_security_group.id]
+    ),
+    opts=pulumi.ResourceOptions(depends_on=[langserve_cluster_instance])
+)
+
 
 # Route53
 hosted_zone = aws.route53.get_zone(name=domain_name)

backend/sql/requirements.txt

@@ -0,0 +1 @@
+psycopg2-binary

backend/sql/sql_migration_lambda_function.py

@@ -0,0 +1,57 @@
+import os
+import psycopg2
+from psycopg2 import sql
+
+def run_migrations():
+    # Database connection parameters
+    db_params = {
+        "host": os.environ['POSTGRES_HOST'],
+        "port": os.environ['POSTGRES_PORT'],
+        "dbname": os.environ['POSTGRES_DB'],
+        "user": os.environ['POSTGRES_USER'],
+        "password": os.environ['POSTGRES_PASSWORD']
+    }
+
+    # Establish a connection to the database
+    conn = psycopg2.connect(**db_params)
+    conn.autocommit = True
+    cursor = conn.cursor()
+
+    # Path to the migration scripts
+    migrations_path = 'code/migrations'
+
+    # Execute each migration script in the migrations directory
+    for root, dirs, files in os.walk(migrations_path):
+        for filename in sorted(files):
+            if filename.endswith('.sql'):
+                print(f"[INFO] Executing {filename}")
+                file_path = os.path.join(root, filename)
+                with open(file_path, 'r') as file:
+                    sql_script = file.read()
+                    print(f"[INFO] Running script: {sql_script}")
+                    try:
+                        cursor.execute(sql.SQL(sql_script))
+                        print(f"[INFO] Executed {filename} successfully.")
+
+                        if cursor.description is not None:
+                            results = cursor.fetchall()
+                            for row in results:
+                                print(row)
+                    except Exception as script_error:
+                        print(f"Error executing {filename}: {script_error}")
+
+    cursor.close()
+    conn.close()
+
+def lambda_handler(event, context):
+    try:
+        run_migrations()
+        return {
+            'statusCode': 200,
+            'body': 'Migrations executed successfully.'
+        }
+    except Exception as e:
+        return {
+            'statusCode': 500,
+            'body': f'Error executing migrations: {str(e)}'
+        }