Skip to content

Commit

Permalink
Move uri generation to update endpoints (#584)
Browse files Browse the repository at this point in the history
  • Loading branch information
dragomirp authored Aug 16, 2024
1 parent da75a5e commit c5e7075
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 16 deletions.
32 changes: 20 additions & 12 deletions src/relations/postgresql_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,6 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
# Share the credentials with the application.
self.database_provides.set_credentials(event.relation.id, user, password)

# Update the read/write and read-only endpoints.
self.update_endpoints(event)

# Set the database version.
self.database_provides.set_version(
event.relation.id, self.charm.postgresql.get_postgresql_version()
Expand All @@ -115,11 +112,8 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
# Set the database name
self.database_provides.set_database(event.relation.id, database)

# Set connection string URI.
self.database_provides.set_uris(
event.relation.id,
f"postgresql://{user}:{password}@{self.charm.primary_endpoint}:{DATABASE_PORT}/{database}",
)
# Update the read/write and read-only endpoints.
self.update_endpoints(event)

self._update_unit_status(event.relation)
except (
Expand Down Expand Up @@ -185,7 +179,11 @@ def update_endpoints(self, event: DatabaseRequestedEvent = None) -> None:

# Get the current relation or all the relations
# if this is triggered by another type of event.
relations = [event.relation] if event else self.model.relations[self.relation_name]
relations_ids = [event.relation.id] if event else None
rel_data = self.database_provides.fetch_relation_data(
relations_ids, ["external-node-connectivity", "database"]
)
secret_data = self.database_provides.fetch_my_relation_data(relations_ids, ["password"])

# If there are no replicas, remove the read-only endpoint.
replicas_endpoint = list(self.charm.members_ips - {self.charm.primary_endpoint})
Expand All @@ -196,19 +194,29 @@ def update_endpoints(self, event: DatabaseRequestedEvent = None) -> None:
else ""
)

for relation in relations:
for relation_id in rel_data.keys():
user = f"relation-{relation_id}"
database = rel_data[relation_id].get("database")
password = secret_data.get(relation_id, {}).get("password")

# Set the read/write endpoint.
self.database_provides.set_endpoints(
relation.id,
relation_id,
f"{self.charm.primary_endpoint}:{DATABASE_PORT}",
)

# Set the read-only endpoint.
self.database_provides.set_read_only_endpoints(
relation.id,
relation_id,
read_only_endpoints,
)

# Set connection string URI.
self.database_provides.set_uris(
relation_id,
f"postgresql://{user}:{password}@{self.charm.primary_endpoint}:{DATABASE_PORT}/{database}",
)

def _check_multiple_endpoints(self) -> bool:
"""Checks if there are relations with other endpoints."""
relation_names = {relation.name for relation in self.charm.client_relations}
Expand Down
43 changes: 39 additions & 4 deletions tests/unit/test_postgresql_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,6 @@ def test_on_database_requested(harness):
"data": f'{{"database": "{DATABASE}", "extra-user-roles": "{EXTRA_USER_ROLES}"}}',
"username": user,
"password": "test-password",
"uris": f"postgresql://{user}:[email protected]:5432/{DATABASE}",
"version": POSTGRESQL_VERSION,
"database": f"{DATABASE}",
}
Expand Down Expand Up @@ -223,7 +222,12 @@ def test_update_endpoints_with_event(harness):
new_callable=PropertyMock,
) as _members_ips,
patch("charm.Patroni.get_primary", new_callable=PropertyMock) as _get_primary,
patch(
"relations.postgresql_provider.DatabaseProvides.fetch_my_relation_data"
) as _fetch_my_relation_data,
):
_fetch_my_relation_data.return_value.get().get.return_value = "test_password"

# Mock the members_ips list to simulate different scenarios
# (with and without a replica).
rel_id = harness.model.get_relation(RELATION_NAME).id
Expand All @@ -232,6 +236,12 @@ def test_update_endpoints_with_event(harness):
# Add two different relations.
rel_id = harness.add_relation(RELATION_NAME, "application")
another_rel_id = harness.add_relation(RELATION_NAME, "application")
with harness.hooks_disabled():
harness.update_relation_data(
rel_id,
"application",
{"database": "test_db", "extra-user-roles": ""},
)

# Define a mock relation changed event to be used in the subsequent update endpoints calls.
mock_event = Mock()
Expand All @@ -246,13 +256,16 @@ def test_update_endpoints_with_event(harness):
assert harness.get_relation_data(rel_id, harness.charm.app.name) == {
"endpoints": "1.1.1.1:5432",
"read-only-endpoints": "2.2.2.2:5432",
"uris": "postgresql://relation-2:[email protected]:5432/test_db",
}
assert harness.get_relation_data(another_rel_id, harness.charm.app.name) == {}
_fetch_my_relation_data.assert_called_once_with([2], ["password"])

# Also test with only a primary instance.
harness.charm.postgresql_client_relation.update_endpoints(mock_event)
assert harness.get_relation_data(rel_id, harness.charm.app.name) == {
"endpoints": "1.1.1.1:5432"
"endpoints": "1.1.1.1:5432",
"uris": "postgresql://relation-2:[email protected]:5432/test_db",
}
assert harness.get_relation_data(another_rel_id, harness.charm.app.name) == {}

Expand All @@ -269,7 +282,12 @@ def test_update_endpoints_without_event(harness):
new_callable=PropertyMock,
) as _members_ips,
patch("charm.Patroni.get_primary", new_callable=PropertyMock) as _get_primary,
patch(
"relations.postgresql_provider.DatabaseProvides.fetch_my_relation_data"
) as _fetch_my_relation_data,
):
_fetch_my_relation_data.return_value.get().get.return_value = "test_password"

# Mock the members_ips list to simulate different scenarios
# (with and without a replica).
_members_ips.side_effect = [{"1.1.1.1", "2.2.2.2"}, {"1.1.1.1"}]
Expand All @@ -279,23 +297,40 @@ def test_update_endpoints_without_event(harness):
rel_id = harness.add_relation(RELATION_NAME, "application")
another_rel_id = harness.add_relation(RELATION_NAME, "application")

with harness.hooks_disabled():
harness.update_relation_data(
rel_id,
"application",
{"database": "test_db", "extra-user-roles": ""},
)
harness.update_relation_data(
another_rel_id,
"application",
{"database": "test_db2", "extra-user-roles": ""},
)

# Test with both a primary and a replica.
# Update the endpoints and check that all relations' databags are updated.
harness.charm.postgresql_client_relation.update_endpoints()
assert harness.get_relation_data(rel_id, harness.charm.app.name) == {
"endpoints": "1.1.1.1:5432",
"read-only-endpoints": "2.2.2.2:5432",
"uris": "postgresql://relation-2:[email protected]:5432/test_db",
}
assert harness.get_relation_data(another_rel_id, harness.charm.app.name) == {
"endpoints": "1.1.1.1:5432",
"read-only-endpoints": "2.2.2.2:5432",
"uris": "postgresql://relation-3:[email protected]:5432/test_db2",
}
_fetch_my_relation_data.assert_called_once_with(None, ["password"])

# Also test with only a primary instance.
harness.charm.postgresql_client_relation.update_endpoints()
assert harness.get_relation_data(rel_id, harness.charm.app.name) == {
"endpoints": "1.1.1.1:5432"
"endpoints": "1.1.1.1:5432",
"uris": "postgresql://relation-2:[email protected]:5432/test_db",
}
assert harness.get_relation_data(another_rel_id, harness.charm.app.name) == {
"endpoints": "1.1.1.1:5432"
"endpoints": "1.1.1.1:5432",
"uris": "postgresql://relation-3:[email protected]:5432/test_db2",
}

0 comments on commit c5e7075

Please sign in to comment.