Skip to content

Commit

Permalink
Set the region
Browse files Browse the repository at this point in the history
  • Loading branch information
WillGibson committed Nov 12, 2024
1 parent 13c5e80 commit 4411efe
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 17 deletions.
12 changes: 6 additions & 6 deletions postgres/manage_users.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,15 @@ 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)
update_db_user_password(conn, cursor, username, password)
else:
create_db_user(conn, cursor, username, password, permissions)
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)
Expand Down Expand Up @@ -73,8 +73,8 @@ def handler(event, context):
username = event['Username']
user_permissions = event['Permissions']

secrets_manager = boto3.client("secretsmanager")
ssm = boto3.client("ssm")
secrets_manager = boto3.client("secretsmanager", region_name="eu-west-2")
ssm = boto3.client("ssm", region_name="eu-west-2")

master_user = json.loads(secrets_manager.get_secret_value(SecretId=db_master_user_secret_arn)["SecretString"])

Expand Down Expand Up @@ -104,7 +104,7 @@ def handler(event, context):
)

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)

Expand Down
22 changes: 11 additions & 11 deletions postgres/tests/test_manage_users.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def test_create_or_update_db_user(self):

conn.commit.assert_called_once()


def test_create_or_update_db_user_when_user_exists(self):
self.cursor.fetchone.return_value = ["test_user"]
conn = MagicMock()
Expand All @@ -74,14 +74,14 @@ def test_create_or_update_db_user_when_user_exists(self):
)

conn.commit.assert_called_once()


@mock_aws
def test_create_or_update_user_secret(self):
ssm = boto3.client("ssm")
ssm = boto3.client("ssm", region_name="eu-west-2")
user_secret_name = "/test/secret"
user_secret_string = {"username": "test_user", "password": "test_password"}

response = create_or_update_user_secret(
ssm, user_secret_name, user_secret_string, self.event
)
Expand All @@ -97,7 +97,7 @@ def test_create_or_update_user_secret(self):

@mock_aws
def test_create_or_update_user_secret_overwrites(self):
ssm = boto3.client("ssm")
ssm = boto3.client("ssm", region_name="eu-west-2")
user_secret_name = "/test/secret"
user_secret_string = {"username": "test_user", "password": "test_password"}
ssm.put_parameter(Name=user_secret_name, Value="blah", Type="String")
Expand All @@ -108,28 +108,28 @@ def test_create_or_update_user_secret_overwrites(self):

assert parameter["Version"] == 2
assert parameter["Value"] == json.dumps(user_secret_string)


@patch("postgres.manage_users.create_or_update_db_user")
@patch("postgres.manage_users.psycopg2.connect")
@mock_aws
def test_handler(self, mock_connect, mock_create_or_update_db_user):
secretsmanager = boto3.client("secretsmanager")
secretsmanager = boto3.client("secretsmanager", region_name="eu-west-2")
secret_id = secretsmanager.create_secret(
Name=self.secret_name, SecretString=self.secret_string
)["ARN"]

self.event["MasterUserSecretArn"] = secret_id

mock_connect.return_value = self.conn
self.conn.cursor = self.cursor

handler(self.event, self.context)

user_password = json.loads(boto3.client("ssm").get_parameter(Name=self.secret_name, WithDecryption=True)["Parameter"]["Value"])[
user_password = json.loads(boto3.client("ssm", region_name="eu-west-2").get_parameter(Name=self.secret_name, WithDecryption=True)["Parameter"]["Value"])[
"password"
]

mock_create_or_update_db_user.assert_called_once_with(
self.conn, self.cursor(), "test-user", user_password, ["SELECT"]
)

0 comments on commit 4411efe

Please sign in to comment.