Dbtp 1298 postgres docker tests

postgres/lambda.tf

@@ -59,8 +59,8 @@ data "aws_security_group" "rds-endpoint" {
 # tflint-ignore: terraform_unused_declarations
 data "archive_file" "lambda" {
   type        = "zip"
-  source_file = "${path.module}/manage_users.py"
-  output_path = "${path.module}/manage_users.zip"
+  source_file = "${path.module}/manage-users/manage_users/manage_users.py"
+  output_path = "${path.module}/manage-users/manage_users/manage_users.zip"
   depends_on = [
     aws_iam_role.lambda-execution-role
   ]

postgres/manage-users/Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.9
+RUN apt-get update && apt-get install -y curl
+RUN curl -sSL https://install.python-poetry.org | python3 -
+ENV PATH="/root/.local/bin:$PATH"
+
+WORKDIR /src
+COPY poetry.lock ./
+COPY pyproject.toml ./
+
+RUN poetry update
+COPY . ./
+RUN poetry install --no-root
+
+HEALTHCHECK --interval=5m --timeout=3s \
+  CMD curl -f http://localhost/ || exit 1
+
+CMD ["poetry", "run", "pytest"]

postgres/manage-users/Makefile

@@ -0,0 +1,8 @@
+build:
+	docker-compose build
+
+test: down build
+	docker-compose run manage-users poetry run pytest
+
+down:
+	docker-compose down

postgres/manage-users/README.md

@@ -0,0 +1,7 @@
+# manage-users
+
+## Setup
+
+```shell
+brew install docker-compose
+```

postgres/manage-users/docker-compose.yml

@@ -0,0 +1,31 @@
+services:
+  manage-users:
+    networks:
+      - testnet
+    build:
+      context: ../manage-users
+    depends_on:
+      - postgres-unittest
+    ports:
+      - '9004:8080'
+    environment:
+      DATABASE_URL: postgres://test_user:test_pw@postgres-unittest:5432/test_db
+    volumes:
+      - ./:/src
+
+  postgres-unittest:
+    networks:
+      - testnet
+    image: postgres:13
+    restart: always
+    command: 'postgres -c log_statement=all -c fsync=off -c vm.zone_reclaim_mode=0 -c synchronous_commit=off -c checkpoint_completion_target=0.9 -c random_page_cost=1.1'
+    ports:
+      - "5432:5432"
+    environment:
+      - POSTGRES_DB=test_db
+      - POSTGRES_USER=test_user
+      - POSTGRES_PASSWORD=test_pw
+
+networks:
+  testnet:
+    driver: bridge

postgres/manage-users/manage_users/manage_users.py

@@ -0,0 +1,108 @@
+import json
+import boto3
+import psycopg2
+from botocore.exceptions import ClientError
+
+
+def create_or_update_db_user(conn, cursor, username, password, permissions):
+    cursor.execute(f"SELECT * FROM pg_catalog.pg_user WHERE usename = '{username}'")
+
+    if cursor.fetchone() is not None:
+      update_db_user_password(conn, cursor, username, password)  
+    else:
+      create_db_user(conn, cursor, username, password, permissions) 
+
+
+def update_db_user_password(conn, cursor, username, password):
+    cursor.execute(f"ALTER USER {username} WITH ENCRYPTED PASSWORD '%s'" % password)
+    conn.commit()
+
+
+def create_db_user(conn, cursor, username, password, permissions):
+    cursor.execute(f"CREATE USER {username} WITH ENCRYPTED PASSWORD '%s'" % password)
+    cursor.execute(f"GRANT {username} to postgres;")
+    cursor.execute(f"GRANT {', '.join(permissions)} ON ALL TABLES IN SCHEMA public TO {username};")
+    cursor.execute(f"ALTER DEFAULT PRIVILEGES FOR USER application_user IN SCHEMA public GRANT {', '.join(permissions)} ON TABLES TO {username};")
+
+    if 'INSERT' in permissions:
+        cursor.execute(f"GRANT CREATE ON SCHEMA public TO {username};")
+
+    conn.commit()
+
+
+def create_or_update_user_secret(ssm, user_secret_name, user_secret_string, event):
+    user_secret_description = event['SecretDescription']
+    copilot_application = event['CopilotApplication']
+    copilot_environment = event['CopilotEnvironment']
+
+    user_secret = None
+
+    try:
+        user_secret = ssm.put_parameter(
+            Name=user_secret_name,
+            Description=user_secret_description,
+            Value=json.dumps(user_secret_string),
+            Tags=[
+                {'Key': 'managed-by', 'Value': 'Terraform'},
+                {'Key': 'copilot-application', 'Value': copilot_application},
+                {'Key': 'copilot-environment', 'Value': copilot_environment},
+            ],
+            Type="SecureString",
+        )
+    except ClientError as error:
+        if error.response["Error"]["Code"] == "ParameterAlreadyExists":
+            user_secret = ssm.put_parameter(
+                Name=user_secret_name,
+                Description=user_secret_description,
+                Value=json.dumps(user_secret_string),
+                Overwrite=True,
+            )
+
+    return user_secret
+
+
+def handler(event, context):
+    print("REQUEST RECEIVED:\n" + json.dumps(event))
+
+    db_master_user_secret_arn = event['MasterUserSecretArn']
+    user_secret_name = event['SecretName']
+    username = event['Username']
+    user_permissions = event['Permissions']
+
+    secrets_manager = boto3.client("secretsmanager")
+    ssm = boto3.client("ssm")
+
+    master_user = json.loads(secrets_manager.get_secret_value(SecretId=db_master_user_secret_arn)["SecretString"])
+
+    user_password = secrets_manager.get_random_password(
+        PasswordLength=16,
+        ExcludeCharacters='[]{}()"@/\\;=?&`><:|#',
+        ExcludePunctuation=True,
+        IncludeSpace=False,
+    )["RandomPassword"]
+
+    user_secret_string = {
+        "username": username,
+        "password": user_password,
+        "engine": event["DbEngine"],
+        "port": event["DbPort"],
+        "dbname": event["DbName"],
+        "host": event["DbHost"],
+        "dbInstanceIdentifier": event["dbInstanceIdentifier"]
+    }
+
+    conn = psycopg2.connect(
+        dbname=event["DbName"],
+        user=master_user["username"],
+        password=master_user["password"],
+        host=event["DbHost"],
+        port=event["DbPort"]
+    )
+
+    cursor = conn.cursor()
+
+    create_or_update_db_user(conn, cursor, username, user_password, user_permissions)
+    create_or_update_user_secret(ssm, user_secret_name, user_secret_string, event)
+
+    cursor.close()
+    conn.close()