Add Athena database user with read-only privileges

jlongland · jlongland · commit 8d72a092d8fb · 2024-08-30T12:39:15.000-07:00
diff --git a/setup/prepare_aurora_db.py b/setup/prepare_aurora_db.py
@@ -2,174 +2,158 @@
 
 import argparse
 import json
-
 import boto3
 from rich.console import Console
 from botocore.exceptions import ClientError
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
     "--stack-name",
-    help="the name of the CloudFormation stack containing the cd2 Aurora database",
+    help="The name of the CloudFormation stack containing the Aurora database",
     required=True,
 )
-
 args = parser.parse_args()
 
 console = Console()
 
+# Initialize AWS clients
 secrets_client = boto3.client("secretsmanager")
 rds_data_client = boto3.client("rds-data")
-cf_reource = boto3.resource("cloudformation")
-stack = cf_reource.Stack(args.stack_name)
+cf_resource = boto3.resource("cloudformation")
+stack = cf_resource.Stack(args.stack_name)
 
 console.print("Starting database preparation", style="bold green")
 
-# read the outputs and parameters from the Cloudformation stack
-stack_outputs = {
-    output["OutputKey"]: output["OutputValue"]
-    for output in stack.outputs
-}
-stack_parameters = {
-    parameter["ParameterKey"]: parameter["ParameterValue"]
-    for parameter in stack.parameters
-}
+# Fetch stack outputs and parameters
+stack_outputs = {output["OutputKey"]: output["OutputValue"] for output in stack.outputs}
+stack_parameters = {parameter["ParameterKey"]: parameter["ParameterValue"] for parameter in stack.parameters}
 
-# get the database admin secret
+# Get admin and cluster details
 admin_secret_arn = stack_outputs["AdminSecretArn"]
-admin_secret = json.loads(
-    secrets_client.get_secret_value(SecretId=admin_secret_arn)["SecretString"]
-)
+admin_secret = json.loads(secrets_client.get_secret_value(SecretId=admin_secret_arn)["SecretString"])
 admin_username = admin_secret["username"]
-
-# get the database cluster ARN
 aurora_cluster_arn = stack_outputs["AuroraClusterArn"]
 
-# get the environment
+# Get environment and resource prefix
 env = stack_parameters["EnvironmentParameter"]
-
-# get the resource prefix
 prefix = stack_parameters["ResourcePrefixParameter"]
 
-# get the database user secrets
-secret_name_prefix = f"{prefix}-cd2-db-user-{env}-"
-user_secrets = secrets_client.list_secrets(
-    Filters=[{"Key": "name", "Values": [secret_name_prefix]}],
-    MaxResults=100,
-)
+# Define role-based privileges
+role_privileges = {
+    "read_only": "GRANT SELECT ON ALL TABLES IN SCHEMA {schema_name} TO {username};",
+    "read_write": "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {schema_name} TO {username};",
+    "admin": "GRANT ALL PRIVILEGES ON SCHEMA {schema_name} TO {username};",
+}
 
-# for each user secret, create the database user and schema
-for s in user_secrets["SecretList"]:
-    secret_arn = s["ARN"]
+# Define user-role mapping
+user_roles = {
+    "athena": "read_only",
+    "canvas": "admin"
+}
 
-    secret_value = json.loads(
-        secrets_client.get_secret_value(SecretId=secret_arn)["SecretString"]
-    )
-    password = secret_value["password"]
-    username = secret_value["username"]
-    database_name = secret_value["dbname"]
-    schema_name = username
+# List of usernames that should have schemas created
+users_to_create_schema = ["canvas"]
 
-    console.print(
-        f" - creating database user [bold]{username}[/bold] in database [bold]{database_name}[/bold]",
-        style="green",
+def get_user_role(username):
+    """Retrieve the role for a given username. Return read-only if not found"""
+    return user_roles.get(username, "read_only")
+
+def execute_statement(sql, database_name):
+    rds_data_client.execute_statement(
+        resourceArn=aurora_cluster_arn,
+        secretArn=admin_secret_arn,
+        sql=sql,
+        database=database_name,
     )
 
-    # create the database user
+def create_user(username, password, database_name):
+    """Create a user"""
     try:
-        user_sql = f"CREATE USER {username} WITH PASSWORD '{password}' LOGIN"
-        rds_data_client.execute_statement(
-            resourceArn=aurora_cluster_arn,
-            secretArn=admin_secret_arn,
-            sql=user_sql,
-            database=database_name,
-        )
-        console.print(" - Created user", style="bold green")
+        create_user_sql = f"CREATE USER {username} WITH PASSWORD '{password}'"
+        execute_statement(create_user_sql, database_name)
+        console.print(f" - Created user {username}", style="bold green")
     except ClientError as e:
         if "already exists" in e.response["Error"]["Message"]:
-            console.print(f" - User {username} already exists", style="bold red")
-
-            try:
-                change_sql = f"ALTER USER {username} WITH PASSWORD '{password}'"
-                rds_data_client.execute_statement(
-                    resourceArn=aurora_cluster_arn,
-                    secretArn=admin_secret_arn,
-                    sql=change_sql,
-                    database=database_name,
-                )
-                console.print(
-                    f" - Updated password for user {username}", style="bold green"
-                )
-            except ClientError as e:
-                console.print(
-                    f" ! Unexpected error when updating password for {username}: {e}", style="bold red"
-                )
-                continue
+            console.print(f" - User {username} already exists. Updating password...", style="yellow")
+            update_password_sql = f"ALTER USER {username} WITH PASSWORD '{password}'"
+            execute_statement(update_password_sql, database_name)
+            console.print(f" - Updated password for user {username}", style="green")
         else:
-            console.print(
-                f" ! Unexpected error when creating user {username}: {e}", style="bold red"
-            )
-            continue
+            console.print(f" ! Error creating user {username}: {e}", style="bold red")
 
-    # Grant the role to the admin user
+def create_schema(schema_name, username, database_name):
+    """Create a schema with user as owner"""
     try:
-        grant_sql = f"GRANT {username} TO {admin_username}"
-        rds_data_client.execute_statement(
-            resourceArn=aurora_cluster_arn,
-            secretArn=admin_secret_arn,
-            sql=grant_sql,
-            database=database_name,
-        )
-        console.print(
-            f" - Granted user {username} to {admin_username}", style="bold green"
-        )
+        create_schema_sql = f"CREATE SCHEMA IF NOT EXISTS {username} AUTHORIZATION {username}"
+        execute_statement(create_schema_sql, database_name)
+        console.print(f" - Created schema {schema_name} with owner {username}", style="bold green")
     except ClientError as e:
-        console.print(f" ! Unexpected error granting {username} role to {admin_username}: {e}", style="bold red")
-        continue
+        if "already exists" in e.response["Error"]["Message"]:
+            console.print(f" - Schema {schema_name} already exists", style="yellow")
+        else:
+            console.print(f" ! Error creating schema {schema_name} with owner {username}: {e}", style="bold red")
 
-    # create the schema
+def assign_privileges(username, schema_name, role, database_name):
+    """Assign privileges to a database user based on their role."""
     try:
-        schema_sql = f"CREATE SCHEMA IF NOT EXISTS AUTHORIZATION {username}"
-        rds_data_client.execute_statement(
-            resourceArn=aurora_cluster_arn,
-            secretArn=admin_secret_arn,
-            sql=schema_sql,
-            database=database_name,
-        )
-        console.print(f" - Created schema [bold]{username}[/bold]", style="green")
+        grant_schema_sql = role_privileges[role].format(username=username, schema_name=schema_name)
+        execute_statement(grant_schema_sql, database_name)
+        console.print(f" - Granted {role} privileges on schema {schema_name} to user {username}", style="bold green")
     except ClientError as e:
-        console.print(f" ! Unexpected error creating schema {username}: {e}", style="bold red")
-        continue
+        console.print(f" ! Error granting {role} privileges on schema {schema_name} to {username}: {e}", style="bold red")
 
-    # create the instructure_dap schema
+def grant_usage_to_schema(username, schema_name, database_name):
+    """Grant usage on a schema to a user"""
     try:
-        schema_sql = f"CREATE SCHEMA IF NOT EXISTS instructure_dap AUTHORIZATION {username}"
-        rds_data_client.execute_statement(
-            resourceArn=aurora_cluster_arn,
-            secretArn=admin_secret_arn,
-            sql=schema_sql,
-            database=database_name,
-        )
-        console.print(
-            f" - Created schema [bold]instructure_dap[/bold] in database [bold]{database_name}[/bold]",
-            style="green",
-        )
+        grant_usage_sql = f"GRANT USAGE ON SCHEMA {schema_name} TO {username}"
+        execute_statement(grant_usage_sql, database_name)
+        console.print(f" - Granted usage on schema {schema_name} to user {username}", style="bold green")
     except ClientError as e:
-        console.print(f" ! Unexpected error: {e}", style="bold red")
-        continue
+        console.print(f" ! Error granting usage on schema {schema_name} to {username}: {e}", style="bold red")
 
-    # grant create permission on database to canvas user
+def grant_user_to_admin(username, admin_username, database_name):
+    """Grant user to the admin user"""
     try:
-        grant_sql = f"GRANT CREATE ON DATABASE {database_name} TO {username}"
-        rds_data_client.execute_statement(
-            resourceArn=aurora_cluster_arn,
-            secretArn=admin_secret_arn,
-            sql=grant_sql,
-            database="postgres",
-        )
-        console.print(
-            f" - Granted CREATE on database {database_name} to user {username}",
-            style="bold green",
-        )
+        grant_user_sql = f"GRANT {username} TO {admin_username}"
+        execute_statement(grant_user_sql, database_name)
+        console.print(f" - Granted user {username} to user {admin_username}", style="bold green")
     except ClientError as e:
-        console.print(f" ! Unexpected error: {e}", style="bold red")
+        console.print(f" ! Error granting user {username} to user {admin_username}: {e}", style="bold red")
+
+# Get all database user secrets
+secret_name_prefix = f"{prefix}-cd2-db-user-{env}-"
+user_secrets = secrets_client.list_secrets(
+    Filters=[{"Key": "name", "Values": [secret_name_prefix]}],
+    MaxResults=100,
+)
+
+# Process each user secret to create database users, schemas, and assign roles
+for s in user_secrets["SecretList"]:
+    secret_arn = s["ARN"]
+    secret_value = json.loads(secrets_client.get_secret_value(SecretId=secret_arn)["SecretString"])
+    username = secret_value["username"]
+    database_name = secret_value["dbname"]
+    
+    # Create or update the user
+    create_user(username, secret_value["password"], database_name)
+    
+    # Grant user to admin user
+    grant_user_to_admin(username, admin_username, database_name)
+    
+    # Create schema for user (with them as owner) if they need a schema
+    if username in users_to_create_schema:
+        create_schema(username, username, database_name)
+    
+    # Create instructure_dap schema for canvas user with them as owner
+    if username == "canvas":
+        create_schema("instructure_dap", username, database_name)
+
+    # Assign privileges to canvas and instructure_dap schemas
+    # Defaults to read-only if user is not set in user_roles dict
+    user_role = get_user_role(username)
+    
+    grant_usage_to_schema(username, "canvas", database_name)
+    assign_privileges(username, "canvas", user_role, database_name)
+    
+    grant_usage_to_schema(username, "instructure_dap", database_name)
+    assign_privileges(username, "instructure_dap", user_role, database_name)
diff --git a/template.yaml b/template.yaml
@@ -1150,7 +1150,7 @@ Resources:
         LambdaFunctionName: !Ref AthenaConnectorLambdaNameParameter
         CompositeHandler: PostGreSqlMuxCompositeHandler
         DefaultConnectionString: !Sub |
-          postgres://jdbc:postgresql://${AuroraDatabaseCluster.Endpoint.Address}:${AuroraDatabaseCluster.Endpoint.Port}/cd2?${!${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-canvas}&sslmode=verify-ca&sslfactory=org.postgresql.ssl.DefaultJavaSSLFactory&stringtype=unspecified
+          postgres://jdbc:postgresql://${AuroraDatabaseCluster.Endpoint.Address}:${AuroraDatabaseCluster.Endpoint.Port}/cd2?${!${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-athena}&sslmode=verify-ca&sslfactory=org.postgresql.ssl.DefaultJavaSSLFactory&stringtype=unspecified
         DefaultScale: '0'
         DisableSpillEncryption: 'false'
         LambdaMemory: '3008'
@@ -1208,7 +1208,7 @@ Resources:
         - Effect: Allow
           Action:
           - secretsmanager:GetSecretValue
-          Resource: !Sub arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-canvas*
+          Resource: !Sub arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-athena*
         - Effect: Allow
           Action:
           - logs:CreateLogGroup
@@ -1260,7 +1260,33 @@ Resources:
         function: !Sub arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${AthenaConnectorLambdaNameParameter}
       Tags: 
         - Key: !Sub ${TagNameParameter}
-          Value: !Sub ${TagValueParameter}          
+          Value: !Sub ${TagValueParameter}
+
+  # Secret for PostgreSQL Athena connector "athena" user
+  DatabaseUserSecretAthena:
+    Type: AWS::SecretsManager::Secret
+    Condition: CreateAthenaConnector
+    Properties:
+      Name: !Sub ${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-athena
+      Description: Database user for Athena PostgreSQL connector
+      GenerateSecretString:
+        SecretStringTemplate: !Sub '{"username": "athena"}'
+        GenerateStringKey: password
+        PasswordLength: 128
+        ExcludePunctuation: true
+      KmsKeyId: !Ref SecretsKmsKey
+      Tags: 
+        - Key: !Sub ${TagNameParameter}
+          Value: !Sub ${TagValueParameter}      
+
+  # Attach Athena database user secret to the cluster
+  DatabaseUserSecretAttachmentAthena:
+    Type: AWS::SecretsManager::SecretTargetAttachment
+    Condition: CreateAthenaConnector
+    Properties:
+      SecretId: !Ref DatabaseUserSecretAthena
+      TargetId: !Ref AuroraDatabaseCluster
+      TargetType: AWS::RDS::DBCluster          
 
 Outputs:
   AdminSecretArn:
