Skip to content

Commit bc9e9db

Browse files
committed
Add Athena database user with read-only privileges
1 parent 847bc7e commit bc9e9db

File tree

2 files changed

+146
-128
lines changed

2 files changed

+146
-128
lines changed

setup/prepare_aurora_db.py

100755100644
Lines changed: 117 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -2,174 +2,166 @@
22

33
import argparse
44
import json
5-
65
import boto3
76
from rich.console import Console
87
from botocore.exceptions import ClientError
98

109
parser = argparse.ArgumentParser()
1110
parser.add_argument(
1211
"--stack-name",
13-
help="the name of the CloudFormation stack containing the cd2 Aurora database",
12+
help="The name of the CloudFormation stack containing the Aurora database",
1413
required=True,
1514
)
16-
1715
args = parser.parse_args()
1816

1917
console = Console()
2018

19+
# Initialize AWS clients
2120
secrets_client = boto3.client("secretsmanager")
2221
rds_data_client = boto3.client("rds-data")
23-
cf_reource = boto3.resource("cloudformation")
24-
stack = cf_reource.Stack(args.stack_name)
22+
cf_resource = boto3.resource("cloudformation")
23+
stack = cf_resource.Stack(args.stack_name)
2524

2625
console.print("Starting database preparation", style="bold green")
2726

28-
# read the outputs and parameters from the Cloudformation stack
29-
stack_outputs = {
30-
output["OutputKey"]: output["OutputValue"]
31-
for output in stack.outputs
32-
}
33-
stack_parameters = {
34-
parameter["ParameterKey"]: parameter["ParameterValue"]
35-
for parameter in stack.parameters
36-
}
27+
# Fetch stack outputs and parameters
28+
stack_outputs = {output["OutputKey"]: output["OutputValue"] for output in stack.outputs}
29+
stack_parameters = {parameter["ParameterKey"]: parameter["ParameterValue"] for parameter in stack.parameters}
3730

38-
# get the database admin secret
31+
# Get admin and cluster details
3932
admin_secret_arn = stack_outputs["AdminSecretArn"]
40-
admin_secret = json.loads(
41-
secrets_client.get_secret_value(SecretId=admin_secret_arn)["SecretString"]
42-
)
33+
admin_secret = json.loads(secrets_client.get_secret_value(SecretId=admin_secret_arn)["SecretString"])
4334
admin_username = admin_secret["username"]
44-
45-
# get the database cluster ARN
4635
aurora_cluster_arn = stack_outputs["AuroraClusterArn"]
4736

48-
# get the environment
37+
# Get environment and resource prefix
4938
env = stack_parameters["EnvironmentParameter"]
50-
51-
# get the resource prefix
5239
prefix = stack_parameters["ResourcePrefixParameter"]
5340

54-
# get the database user secrets
55-
secret_name_prefix = f"{prefix}-cd2-db-user-{env}-"
56-
user_secrets = secrets_client.list_secrets(
57-
Filters=[{"Key": "name", "Values": [secret_name_prefix]}],
58-
MaxResults=100,
59-
)
41+
# Define role-based privileges
42+
role_privileges = {
43+
"read_only": "GRANT SELECT ON ALL TABLES IN SCHEMA {schema_name} TO {username};",
44+
"read_write": "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {schema_name} TO {username};",
45+
"admin": "GRANT ALL PRIVILEGES ON SCHEMA {schema_name} TO {username};",
46+
}
6047

61-
# for each user secret, create the database user and schema
62-
for s in user_secrets["SecretList"]:
63-
secret_arn = s["ARN"]
48+
# Define user-role mapping
49+
user_roles = {
50+
"athena": "read_only",
51+
"canvas": "admin"
52+
}
6453

65-
secret_value = json.loads(
66-
secrets_client.get_secret_value(SecretId=secret_arn)["SecretString"]
67-
)
68-
password = secret_value["password"]
69-
username = secret_value["username"]
70-
database_name = secret_value["dbname"]
71-
schema_name = username
54+
# List of usernames that should have schemas created
55+
users_to_create_schema = ["canvas"]
56+
57+
def get_user_role(username):
58+
"""Retrieve the role for a given username. Return read-only if not found"""
59+
return user_roles.get(username, "read_only")
7260

73-
console.print(
74-
f" - creating database user [bold]{username}[/bold] in database [bold]{database_name}[/bold]",
75-
style="green",
61+
def execute_statement(sql, database_name):
62+
rds_data_client.execute_statement(
63+
resourceArn=aurora_cluster_arn,
64+
secretArn=admin_secret_arn,
65+
sql=sql,
66+
database=database_name,
7667
)
7768

78-
# create the database user
69+
def create_user(username, password, database_name):
70+
"""Create a user"""
7971
try:
80-
user_sql = f"CREATE USER {username} WITH PASSWORD '{password}' LOGIN"
81-
rds_data_client.execute_statement(
82-
resourceArn=aurora_cluster_arn,
83-
secretArn=admin_secret_arn,
84-
sql=user_sql,
85-
database=database_name,
86-
)
87-
console.print(" - Created user", style="bold green")
72+
create_user_sql = f"CREATE USER {username} WITH PASSWORD '{password}'"
73+
execute_statement(create_user_sql, database_name)
74+
console.print(f" - Created user {username}", style="bold green")
8875
except ClientError as e:
8976
if "already exists" in e.response["Error"]["Message"]:
90-
console.print(f" - User {username} already exists", style="bold red")
91-
92-
try:
93-
change_sql = f"ALTER USER {username} WITH PASSWORD '{password}'"
94-
rds_data_client.execute_statement(
95-
resourceArn=aurora_cluster_arn,
96-
secretArn=admin_secret_arn,
97-
sql=change_sql,
98-
database=database_name,
99-
)
100-
console.print(
101-
f" - Updated password for user {username}", style="bold green"
102-
)
103-
except ClientError as e:
104-
console.print(
105-
f" ! Unexpected error when updating password for {username}: {e}", style="bold red"
106-
)
107-
continue
77+
console.print(f" - User {username} already exists. Updating password...", style="yellow")
78+
update_password_sql = f"ALTER USER {username} WITH PASSWORD '{password}'"
79+
execute_statement(update_password_sql, database_name)
80+
console.print(f" - Updated password for user {username}", style="green")
10881
else:
109-
console.print(
110-
f" ! Unexpected error when creating user {username}: {e}", style="bold red"
111-
)
112-
continue
82+
console.print(f" ! Error creating user {username}: {e}", style="bold red")
11383

114-
# Grant the role to the admin user
84+
def create_schema(schema_name, username, database_name):
85+
"""Create a schema with user as owner"""
11586
try:
116-
grant_sql = f"GRANT {username} TO {admin_username}"
117-
rds_data_client.execute_statement(
118-
resourceArn=aurora_cluster_arn,
119-
secretArn=admin_secret_arn,
120-
sql=grant_sql,
121-
database=database_name,
122-
)
123-
console.print(
124-
f" - Granted user {username} to {admin_username}", style="bold green"
125-
)
87+
create_schema_sql = f"CREATE SCHEMA IF NOT EXISTS {username} AUTHORIZATION {username}"
88+
execute_statement(create_schema_sql, database_name)
89+
console.print(f" - Created schema {schema_name} with owner {username}", style="bold green")
12690
except ClientError as e:
127-
console.print(f" ! Unexpected error granting {username} role to {admin_username}: {e}", style="bold red")
128-
continue
129-
130-
# create the schema
91+
if "already exists" in e.response["Error"]["Message"]:
92+
console.print(f" - Schema {schema_name} already exists", style="yellow")
93+
else:
94+
console.print(f" ! Error creating schema {schema_name} with owner {username}: {e}", style="bold red")
95+
96+
def generate_privilege_statements(schema_name, user_roles, role_privileges):
97+
statements = []
98+
for username, role in user_roles.items():
99+
if role in role_privileges:
100+
statement = role_privileges[role].format(schema_name=schema_name, username=username)
101+
statements.append(statement)
102+
return statements
103+
104+
def assign_privileges(username, schema_name, role, database_name):
105+
"""Assign privileges to a database user based on their role."""
131106
try:
132-
schema_sql = f"CREATE SCHEMA IF NOT EXISTS AUTHORIZATION {username}"
133-
rds_data_client.execute_statement(
134-
resourceArn=aurora_cluster_arn,
135-
secretArn=admin_secret_arn,
136-
sql=schema_sql,
137-
database=database_name,
138-
)
139-
console.print(f" - Created schema [bold]{username}[/bold]", style="green")
107+
grant_schema_sql = role_privileges[role].format(username=username, schema_name=schema_name)
108+
execute_statement(grant_schema_sql, database_name)
109+
console.print(f" - Granted {role} privileges on schema {schema_name} to user {username}", style="bold green")
140110
except ClientError as e:
141-
console.print(f" ! Unexpected error creating schema {username}: {e}", style="bold red")
142-
continue
111+
console.print(f" ! Error granting {role} privileges on schema {schema_name} to {username}: {e}", style="bold red")
143112

144-
# create the instructure_dap schema
113+
def grant_usage_to_schema(username, schema_name, database_name):
114+
"""Grant usage on a schema to a user"""
145115
try:
146-
schema_sql = f"CREATE SCHEMA IF NOT EXISTS instructure_dap AUTHORIZATION {username}"
147-
rds_data_client.execute_statement(
148-
resourceArn=aurora_cluster_arn,
149-
secretArn=admin_secret_arn,
150-
sql=schema_sql,
151-
database=database_name,
152-
)
153-
console.print(
154-
f" - Created schema [bold]instructure_dap[/bold] in database [bold]{database_name}[/bold]",
155-
style="green",
156-
)
116+
grant_usage_sql = f"GRANT USAGE ON SCHEMA {schema_name} TO {username}"
117+
execute_statement(grant_usage_sql, database_name)
118+
console.print(f" - Granted usage on schema {schema_name} to user {username}", style="bold green")
157119
except ClientError as e:
158-
console.print(f" ! Unexpected error: {e}", style="bold red")
159-
continue
120+
console.print(f" ! Error granting usage on schema {schema_name} to {username}: {e}", style="bold red")
160121

161-
# grant create permission on database to canvas user
122+
def grant_user_to_admin(username, admin_username, database_name):
123+
"""Grant user to the admin user"""
162124
try:
163-
grant_sql = f"GRANT CREATE ON DATABASE {database_name} TO {username}"
164-
rds_data_client.execute_statement(
165-
resourceArn=aurora_cluster_arn,
166-
secretArn=admin_secret_arn,
167-
sql=grant_sql,
168-
database="postgres",
169-
)
170-
console.print(
171-
f" - Granted CREATE on database {database_name} to user {username}",
172-
style="bold green",
173-
)
125+
grant_user_sql = f"GRANT {username} TO {admin_username}"
126+
execute_statement(grant_user_sql, database_name)
127+
console.print(f" - Granted user {username} to user {admin_username}", style="bold green")
174128
except ClientError as e:
175-
console.print(f" ! Unexpected error: {e}", style="bold red")
129+
console.print(f" ! Error granting user {username} to user {admin_username}: {e}", style="bold red")
130+
131+
# Get all database user secrets
132+
secret_name_prefix = f"{prefix}-cd2-db-user-{env}-"
133+
user_secrets = secrets_client.list_secrets(
134+
Filters=[{"Key": "name", "Values": [secret_name_prefix]}],
135+
MaxResults=100,
136+
)
137+
138+
# Process each user secret to create database users, schemas, and assign roles
139+
for s in user_secrets["SecretList"]:
140+
secret_arn = s["ARN"]
141+
secret_value = json.loads(secrets_client.get_secret_value(SecretId=secret_arn)["SecretString"])
142+
username = secret_value["username"]
143+
database_name = secret_value["dbname"]
144+
145+
# Create or update the user
146+
create_user(username, secret_value["password"], database_name)
147+
148+
# Grant user to admin user
149+
grant_user_to_admin(username, admin_username, database_name)
150+
151+
# Create schema for user (with them as owner) if they need a schema
152+
if username in users_to_create_schema:
153+
create_schema(username, username, database_name)
154+
155+
# Create instructure_dap schema for canvas user with them as owner
156+
if username == "canvas":
157+
create_schema("instructure_dap", username, database_name)
158+
159+
# Assign privileges to canvas and instructure_dap schemas
160+
# Defaults to read-only if user is not set in user_roles dict
161+
user_role = get_user_role(username)
162+
163+
grant_usage_to_schema(username, "canvas", database_name)
164+
assign_privileges(username, "canvas", user_role, database_name)
165+
166+
grant_usage_to_schema(username, "instructure_dap", database_name)
167+
assign_privileges(username, "instructure_dap", user_role, database_name)

template.yaml

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1150,7 +1150,7 @@ Resources:
11501150
LambdaFunctionName: !Ref AthenaConnectorLambdaNameParameter
11511151
CompositeHandler: PostGreSqlMuxCompositeHandler
11521152
DefaultConnectionString: !Sub |
1153-
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
1153+
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
11541154
DefaultScale: '0'
11551155
DisableSpillEncryption: 'false'
11561156
LambdaMemory: '3008'
@@ -1208,7 +1208,7 @@ Resources:
12081208
- Effect: Allow
12091209
Action:
12101210
- secretsmanager:GetSecretValue
1211-
Resource: !Sub arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-canvas*
1211+
Resource: !Sub arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-athena*
12121212
- Effect: Allow
12131213
Action:
12141214
- logs:CreateLogGroup
@@ -1260,7 +1260,33 @@ Resources:
12601260
function: !Sub arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${AthenaConnectorLambdaNameParameter}
12611261
Tags:
12621262
- Key: !Sub ${TagNameParameter}
1263-
Value: !Sub ${TagValueParameter}
1263+
Value: !Sub ${TagValueParameter}
1264+
1265+
# Secret for PostgreSQL Athena connector "athena" user
1266+
DatabaseUserSecretAthena:
1267+
Type: AWS::SecretsManager::Secret
1268+
Condition: CreateAthenaConnector
1269+
Properties:
1270+
Name: !Sub ${ResourcePrefixParameter}-cd2-db-user-${EnvironmentParameter}-athena
1271+
Description: Database user for Athena PostgreSQL connector
1272+
GenerateSecretString:
1273+
SecretStringTemplate: !Sub '{"username": "athena"}'
1274+
GenerateStringKey: password
1275+
PasswordLength: 128
1276+
ExcludePunctuation: true
1277+
KmsKeyId: !Ref SecretsKmsKey
1278+
Tags:
1279+
- Key: !Sub ${TagNameParameter}
1280+
Value: !Sub ${TagValueParameter}
1281+
1282+
# Attach Athena database user secret to the cluster
1283+
DatabaseUserSecretAttachmentAthena:
1284+
Type: AWS::SecretsManager::SecretTargetAttachment
1285+
Condition: CreateAthenaConnector
1286+
Properties:
1287+
SecretId: !Ref DatabaseUserSecretAthena
1288+
TargetId: !Ref AuroraDatabaseCluster
1289+
TargetType: AWS::RDS::DBCluster
12641290

12651291
Outputs:
12661292
AdminSecretArn:

0 commit comments

Comments
 (0)