Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Allow communication between sessions of same user #1708

Merged
merged 1 commit into from
Aug 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions backend/capellacollab/sessions/hooks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
http,
interface,
jupyter,
networking,
persistent_workspace,
provisioning,
pure_variants,
Expand All @@ -29,6 +30,7 @@
"read_only_hook": read_only_workspace.ReadOnlyWorkspaceHook(),
"provisioning": provisioning.ProvisionWorkspaceHook(),
"session_preparation": session_preparation.GitRepositoryCloningHook(),
"networking": networking.NetworkingIntegration(),
}


Expand Down
38 changes: 38 additions & 0 deletions backend/capellacollab/sessions/hooks/networking.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
# SPDX-FileCopyrightText: Copyright DB InfraGO AG and contributors
# SPDX-License-Identifier: Apache-2.0


from capellacollab.sessions import operators
from capellacollab.users import models as users_models

from .. import models as sessions_models
from . import interface


class NetworkingIntegration(interface.HookRegistration):
"""Allow sessions of the same user to talk to each other."""

def post_session_creation_hook( # type: ignore
self,
session_id: str,
operator: operators.KubernetesOperator,
user: users_models.DatabaseUser,
**kwargs,
) -> interface.PostSessionCreationHookResult:
"""Allow sessions of the user to talk to each other."""

operator.create_network_policy_from_pod_to_label(
session_id,
{"capellacollab/session-id": session_id},
{"capellacollab/owner-id": str(user.id)},
)

return interface.PostSessionCreationHookResult()

def pre_session_termination_hook( # type: ignore
self,
operator: operators.KubernetesOperator,
session: sessions_models.DatabaseSession,
**kwargs,
):
operator.delete_network_policy(session.id)
50 changes: 49 additions & 1 deletion backend/capellacollab/sessions/operators/k8s.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@
volumes: list[models.Volume],
init_volumes: list[models.Volume],
annotations: dict[str, str],
labels: dict[str, str],
prometheus_path="/metrics",
prometheus_port=9118,
) -> Session:
Expand All @@ -125,6 +126,7 @@
init_volumes=init_volumes,
tool_resources=tool.config.resources,
annotations=annotations,
labels=labels,
)

self._create_disruption_budget(
Expand Down Expand Up @@ -440,6 +442,51 @@
f"The Kubernetes operator encountered an unsupported session volume type '{type(volume)}'"
)

def create_network_policy_from_pod_to_label(
self,
name: str,
match_labels_from: dict[str, str],
match_labels_to: dict[str, str],
):
return self.v1_networking.create_namespaced_network_policy(
namespace,
client.V1NetworkPolicy(
kind="NetworkPolicy",
api_version="networking.k8s.io/v1",
metadata=client.V1ObjectMeta(name=name),
spec=client.V1NetworkPolicySpec(
pod_selector=client.V1LabelSelector(
match_labels=match_labels_to
),
policy_types=["Ingress"],
ingress=[
client.V1NetworkPolicyIngressRule(
_from=[
client.V1NetworkPolicyPeer(
pod_selector=client.V1LabelSelector(
match_labels=match_labels_from
)
)
],
)
],
),
),
)

def delete_network_policy(self, name: str):
try:
self.v1_networking.delete_namespaced_network_policy(
name,
namespace,
)
except exceptions.ApiException as e:

Check warning on line 483 in backend/capellacollab/sessions/operators/k8s.py

View check run for this annotation

Codecov / codecov/patch

backend/capellacollab/sessions/operators/k8s.py#L483

Added line #L483 was not covered by tests
# Network policy doesn't exist or was already deleted
# Nothing to do
if e.status == http.HTTPStatus.NOT_FOUND:
return
raise

Check warning on line 488 in backend/capellacollab/sessions/operators/k8s.py

View check run for this annotation

Codecov / codecov/patch

backend/capellacollab/sessions/operators/k8s.py#L487-L488

Added lines #L487 - L488 were not covered by tests

def _create_deployment(
self,
image: str,
Expand All @@ -451,6 +498,7 @@
init_volumes: list[models.Volume],
tool_resources: tools_models.Resources,
annotations: dict[str, str],
labels: dict[str, str],
) -> client.V1Deployment:
k8s_volumes, k8s_volume_mounts = self._map_volumes_to_k8s_volumes(
volumes
Expand Down Expand Up @@ -541,7 +589,7 @@
selector=client.V1LabelSelector(match_labels={"app": name}),
template=client.V1PodTemplateSpec(
metadata=client.V1ObjectMeta(
labels={"app": name, "workload": "session"},
labels={"app": name, "workload": "session"} | labels,
annotations=annotations,
),
spec=client.V1PodSpec(
Expand Down
6 changes: 6 additions & 0 deletions backend/capellacollab/sessions/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,11 @@ def request_session(
"capellacollab/connection-method-name": connection_method.name,
}

labels: dict[str, str] = {
"capellacollab/session-id": session_id,
"capellacollab/owner-id": str(user.id),
}

session = operator.start_session(
session_id=session_id,
image=docker_image,
Expand All @@ -191,6 +196,7 @@ def request_session(
volumes=volumes,
init_volumes=init_volumes,
annotations=annotations,
labels=labels,
prometheus_path=tool.config.monitoring.prometheus.path,
prometheus_port=connection_method.ports.metrics,
)
Expand Down
66 changes: 66 additions & 0 deletions backend/tests/sessions/hooks/test_networking_hook.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# SPDX-FileCopyrightText: Copyright DB InfraGO AG and contributors
# SPDX-License-Identifier: Apache-2.0


import kubernetes.client
import pytest

from capellacollab.sessions import models as sessions_models
from capellacollab.sessions import operators
from capellacollab.sessions.hooks import networking as networking_hook
from capellacollab.users import models as users_models


def test_network_policy_created(
user: users_models.DatabaseUser, monkeypatch: pytest.MonkeyPatch
):
network_policy_counter = 0

def mock_create_namespaced_network_policy(
self,
namespace: str,
network_policy: kubernetes.client.V1PersistentVolumeClaim,
):
nonlocal network_policy_counter
network_policy_counter += 1

monkeypatch.setattr(
kubernetes.client.NetworkingV1Api,
"create_namespaced_network_policy",
mock_create_namespaced_network_policy,
)

networking_hook.NetworkingIntegration().post_session_creation_hook(
session_id="test",
operator=operators.KubernetesOperator(),
user=user,
)

assert network_policy_counter == 1


def test_network_policy_deleted(
session: sessions_models.DatabaseSession, monkeypatch: pytest.MonkeyPatch
):
network_policy_del_counter = 0

def mock_delete_namespaced_network_policy(
self,
name: str,
namespace: str,
):
nonlocal network_policy_del_counter
network_policy_del_counter += 1

monkeypatch.setattr(
kubernetes.client.NetworkingV1Api,
"delete_namespaced_network_policy",
mock_delete_namespaced_network_policy,
)

networking_hook.NetworkingIntegration().pre_session_termination_hook(
operator=operators.KubernetesOperator(),
session=session,
)

assert network_policy_del_counter == 1
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ def create_namespaced_pod_disruption_budget(namespace, budget):
ports={"rdp": 3389},
volumes=[],
init_volumes=[],
labels={},
annotations={},
)

Expand Down
2 changes: 1 addition & 1 deletion helm/templates/backend/backend.serviceaccount.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ rules:
resources: ["poddisruptionbudgets"]
verbs: ["create", "delete"]
- apiGroups: ["networking.k8s.io"]
resources: ["ingresses"]
resources: ["ingresses", "networkpolicies"]
verbs: ["create", "delete"]
---
apiVersion: rbac.authorization.k8s.io/v1
Expand Down
Loading