feat: enable toggling of hub log forwarding

src/api/api.py

@@ -6,6 +6,7 @@
 from fastapi.responses import JSONResponse
 
 from src.utils.hub_client import init_hub_client_with_robot, get_node_id_by_robot
+from src.utils.other import extract_hub_envs
 from src.api.oauth import valid_access_token
 from src.resources.database.entity import Database
 from src.resources.analysis.entity import CreateAnalysis
@@ -24,13 +25,9 @@
 class PodOrchestrationAPI:
     def __init__(self, database: Database, namespace: str = 'default'):
         self.database = database
-        robot_id, robot_secret, hub_url_core, hub_auth, http_proxy, https_proxy = (os.getenv('HUB_ROBOT_USER'),
-                                                                                   os.getenv('HUB_ROBOT_SECRET'),
-                                                                                   os.getenv('HUB_URL_CORE'),
-                                                                                   os.getenv('HUB_URL_AUTH'),
-                                                                                   os.getenv('PO_HTTP_PROXY'),
-                                                                                   os.getenv('PO_HTTPS_PROXY'))
+        robot_id, robot_secret, hub_url_core, hub_auth, enable_hub_logging, http_proxy, https_proxy = extract_hub_envs()
 
+        self.enable_hub_logging = enable_hub_logging
         self.hub_core_client = init_hub_client_with_robot(robot_id,
                                                           robot_secret,
                                                           hub_url_core,
@@ -229,7 +226,7 @@ def cleanup_call(self, cleanup_type: str):
 
     def stream_logs_call(self, body: CreateLogEntity):
         try:
-            return stream_logs(body, self.node_id, self.database, self.hub_core_client)
+            return stream_logs(body, self.node_id, self.enable_hub_logging, self.database, self.hub_core_client)
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error streaming logs: {e}")
 
@@ -238,4 +235,4 @@ def health_call(self):
         if not main_alive:
             raise RuntimeError("Main thread is not alive.")
         else:
-            return {"status": "ok"}
+            return {'status': "ok"}

src/api/oauth.py

@@ -7,20 +7,20 @@
 from typing import Annotated
 
 oauth2_scheme = OAuth2AuthorizationCodeBearer(
-    tokenUrl=os.getenv("KEYCLOAK_URL") + "/realms/flame/protocol/openid-connect/token",
-    authorizationUrl=os.getenv("KEYCLOAK_URL") + "/realms/flame/protocol/openid-connect/auth",
-    refreshUrl=os.getenv("KEYCLOAK_URL") + "/realms/flame/protocol/openid-connect/token",
+    tokenUrl=os.getenv('KEYCLOAK_URL') + "/realms/flame/protocol/openid-connect/token",
+    authorizationUrl=os.getenv('KEYCLOAK_URL') + "/realms/flame/protocol/openid-connect/auth",
+    refreshUrl=os.getenv('KEYCLOAK_URL') + "/realms/flame/protocol/openid-connect/token",
 )
 
 
 async def valid_access_token(token: Annotated[str, Depends(oauth2_scheme)]) -> dict:
-    url = os.getenv("KEYCLOAK_URL") + "/realms/flame/protocol/openid-connect/certs"
+    url = os.getenv('KEYCLOAK_URL') + "/realms/flame/protocol/openid-connect/certs"
     jwks_client = PyJWKClient(url)
 
     try:
         sig_key = jwks_client.get_signing_key_from_jwt(token)
         return jwt.decode(token,
                           key=sig_key,
-                          options={"verify_signature": True, "verify_aud": False, "exp": True})
+                          options={'verify_signature': True, 'verify_aud': False, 'exp': True})
     except jwt.exceptions.InvalidTokenError:
         raise HTTPException(status_code=401, detail="Not authenticated")

src/k8s/utils.py

@@ -1,16 +1,32 @@
+import time
 from typing import Literal, Optional, Union
 
-from kubernetes import client
+from kubernetes import config, client
 
 
-def get_k8s_resource_names(resource_type: str,
-                           selector_type: Optional[Literal['label', 'field']] = None,
-                           selector_arg: Optional[str] = None,
-                           manual_name_selector: Optional[str] = None,
-                           namespace: str = "default") -> Optional[Union[str, list[str]]]:
-    if resource_type not in ['deployment', 'pod', 'service', 'networkpolicy', 'configmap']:
+def load_cluster_config():
+    config.load_incluster_config()
+
+
+def get_current_namespace() -> str:
+    namespace_file = '/var/run/secrets/kubernetes.io/serviceaccount/namespace'
+    try:
+        with open(namespace_file, 'r') as file:
+            return file.read().strip()
+    # Handle the case where the file is not found
+    except FileNotFoundError:
+        # Fallback to a default namespace if the file is not found
+        return 'default'
+
+
+def find_k8s_resources(resource_type: str,
+                       selector_type: Optional[Literal['label', 'field']] = None,
+                       selector_arg: Optional[str] = None,
+                       manual_name_selector: Optional[str] = None,
+                       namespace: str = "default") -> Optional[Union[str, list[str]]]:
+    if resource_type not in ['deployment', 'pod', 'service', 'networkpolicy', 'configmap', 'job']:
         raise ValueError("For k8s resource search: resource_type must be one of 'deployment', 'pod', 'service', "
-                         "'networkpolicy', or 'configmap'")
+                         "'networkpolicy', 'configmap', or 'job")
     if (selector_type is not None) and (selector_type not in ['label', 'field']):
         raise ValueError("For k8s resource search: selector_type must be either 'label' or 'field'")
     if (selector_type is not None) and (selector_arg is None):
@@ -19,7 +35,7 @@ def get_k8s_resource_names(resource_type: str,
     kwargs = {'namespace': namespace}
     if selector_type:
         kwargs[f'{selector_type}_selector'] = selector_arg
-        
+
     if resource_type == 'deployment':
         resources = client.AppsV1Api().list_namespaced_deployment(**kwargs)
     elif resource_type == 'networkpolicy':
@@ -32,6 +48,8 @@ def get_k8s_resource_names(resource_type: str,
             resources = core_client.list_namespaced_service(**kwargs)
         elif resource_type == 'configmap':
             resources = core_client.list_namespaced_config_map(**kwargs)
+    elif resource_type == 'job':
+        resources = client.BatchV1Api().list_namespaced_job(**kwargs)
     else:
         raise ValueError(f"Uncaptured resource type discovered! Message the Devs... (found={resource_type})")
 
@@ -52,12 +70,55 @@ def get_k8s_resource_names(resource_type: str,
                 return None
 
 
-def get_current_namespace() -> str:
-    namespace_file = '/var/run/secrets/kubernetes.io/serviceaccount/namespace'
-    try:
-        with open(namespace_file, 'r') as file:
-            return file.read().strip()
-    # Handle the case where the file is not found
-    except FileNotFoundError:
-        # Fallback to a default namespace if the file is not found
-        return 'default'
+def delete_k8s_resource(name: str, resource_type: str, namespace: str = 'default') -> None:
+    """
+    Deletes a Kubernetes resource by name and type.
+    :param name: Name of the resource to delete.
+    :param resource_type: Type of the resource (e.g., 'deployment', 'service', 'pod', 'configmap', 'job').
+    :param namespace: Namespace in which the resource exists.
+    """
+    print(f"PO ACTION - Deleting resource: {name} of type {resource_type} in namespace {namespace} at {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    if resource_type == 'deployment':
+        try:
+            app_client = client.AppsV1Api()
+            app_client.delete_namespaced_deployment(name=name, namespace=namespace, propagation_policy='Foreground')
+        except client.exceptions.ApiException as e:
+            if e.reason != 'Not Found':
+                print(f"Error: Not Found {name} deployment")
+    elif resource_type == 'service':
+        try:
+            core_client = client.CoreV1Api()
+            core_client.delete_namespaced_service(name=name, namespace=namespace)
+        except client.exceptions.ApiException as e:
+            if e.reason != 'Not Found':
+                print(f"Error: Not Found {name} service")
+    elif resource_type == 'pod':
+        try:
+            core_client = client.CoreV1Api()
+            core_client.delete_namespaced_pod(name=name, namespace=namespace)
+        except client.exceptions.ApiException as e:
+            if e.reason != 'Not Found':
+                print(f"Error: Not Found {name} pod")
+    elif resource_type == 'configmap':
+        try:
+            core_client = client.CoreV1Api()
+            core_client.delete_namespaced_config_map(name=name, namespace=namespace)
+        except client.exceptions.ApiException as e:
+            if e.reason != 'Not Found':
+                print(f"Error: Not Found {name} configmap")
+    elif resource_type == 'networkpolicy':
+        try:
+            network_client = client.NetworkingV1Api()
+            network_client.delete_namespaced_network_policy(name=name, namespace=namespace)
+        except client.exceptions.ApiException as e:
+            if e.reason != 'Not Found':
+                print(f"Error: Not Found {name} networkpolicy")
+    elif resource_type == 'job':
+        try:
+            batch_client = client.BatchV1Api()
+            batch_client.delete_namespaced_job(name=name, namespace=namespace, propagation_policy='Foreground')
+        except client.exceptions.ApiException as e:
+            if e.reason != 'Not Found':
+                print(f"Error: Not Found {name} job")
+    else:
+        raise ValueError(f"Unsupported resource type: {resource_type}")

src/main.py

@@ -5,8 +5,7 @@
 
 from src.resources.database.entity import Database
 from src.api.api import PodOrchestrationAPI
-from src.k8s.kubernetes import load_cluster_config
-from src.k8s.utils import get_current_namespace
+from src.k8s.utils import get_current_namespace, load_cluster_config
 from src.status.status import status_loop
 
 

src/resources/database/entity.py

@@ -11,12 +11,12 @@
 
 class Database:
     def __init__(self) -> None:
-        host = os.getenv("POSTGRES_HOST")
+        host = os.getenv('POSTGRES_HOST')
         port = "5432"
-        user = os.getenv("POSTGRES_USER")
-        password = os.getenv("POSTGRES_PASSWORD")
-        database = os.getenv("POSTGRES_DB")
-        conn_uri = f'postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}'
+        user = os.getenv('POSTGRES_USER')
+        password = os.getenv('POSTGRES_PASSWORD')
+        database = os.getenv('POSTGRES_DB')
+        conn_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
         print(conn_uri)
         self.engine = create_engine(conn_uri,
                                     pool_pre_ping=True,
@@ -31,11 +31,11 @@ def reset_db(self) -> None:
 
     def get_deployment(self, deployment_name: str) -> Optional[AnalysisDB]:
         with self.SessionLocal() as session:
-            return session.query(AnalysisDB).filter_by(**{"deployment_name": deployment_name}).first()
+            return session.query(AnalysisDB).filter_by(**{'deployment_name': deployment_name}).first()
 
     def get_latest_deployment(self, analysis_id: str) -> Optional[AnalysisDB]:
         with self.SessionLocal() as session:
-            deployments = session.query(AnalysisDB).filter_by(**{"analysis_id": analysis_id}).all()
+            deployments = session.query(AnalysisDB).filter_by(**{'analysis_id': analysis_id}).all()
             if deployments:
                 return deployments[-1]
             return None
@@ -50,7 +50,7 @@ def analysis_is_running(self, analysis_id: str) -> bool:
 
     def get_deployments(self, analysis_id: str) -> list[AnalysisDB]:
         with self.SessionLocal() as session:
-            return session.query(AnalysisDB).filter_by(**{"analysis_id": analysis_id}).all()
+            return session.query(AnalysisDB).filter_by(**{'analysis_id': analysis_id}).all()
 
     def create_analysis(self,
                         analysis_id: str,
@@ -86,7 +86,7 @@ def create_analysis(self,
 
     def update_analysis(self, analysis_id: str, **kwargs) -> list[AnalysisDB]:
         with self.SessionLocal() as session:
-            analysis = session.query(AnalysisDB).filter_by(**{"analysis_id": analysis_id}).all()
+            analysis = session.query(AnalysisDB).filter_by(**{'analysis_id': analysis_id}).all()
             for deployment in analysis:
                 if deployment:
                     for key, value in kwargs.items():
@@ -97,15 +97,15 @@ def update_analysis(self, analysis_id: str, **kwargs) -> list[AnalysisDB]:
 
     def update_deployment(self, deployment_name: str, **kwargs) -> AnalysisDB:
         with self.SessionLocal() as session:
-            deployment = session.query(AnalysisDB).filter_by(**{"deployment_name": deployment_name}).first()
+            deployment = session.query(AnalysisDB).filter_by(**{'deployment_name': deployment_name}).first()
             for key, value in kwargs.items():
                 setattr(deployment, key, value)
             session.commit()
             return deployment
 
     def delete_analysis(self, analysis_id: str) -> None:
         with self.SessionLocal() as session:
-            analysis = session.query(AnalysisDB).filter_by(**{"analysis_id": analysis_id}).all()
+            analysis = session.query(AnalysisDB).filter_by(**{'analysis_id': analysis_id}).all()
             for deployment in analysis:
                 if deployment:
                     session.delete(deployment)
@@ -154,7 +154,7 @@ def update_analysis_status(self, analysis_id: str, status: str) -> None:
         self.update_analysis(analysis_id, status=status)
 
     def update_deployment_status(self, deployment_name: str, status: str) -> None:
-        print(f"Updating deployment {deployment_name} to status {status}")
+        print(f"PO ACTION - Updating deployment {deployment_name} to status {status}")
         self.update_deployment(deployment_name, status=status)
 
     def stop_analysis(self, analysis_id: str) -> None:
@@ -164,19 +164,19 @@ def extract_analysis_body(self, analysis_id: str) -> Optional[dict]:
         analysis = self.get_deployments(analysis_id)
         if analysis:
             analysis = analysis[0]
-            return {"analysis_id": analysis.analysis_id,
-                    "project_id": analysis.project_id,
-                    "registry_url": analysis.registry_url,
-                    "image_url": analysis.image_url,
-                    "registry_user": analysis.registry_user,
-                    "registry_password": analysis.registry_password,
-                    "namespace": analysis.namespace,
-                    "kong_token": analysis.kong_token,
-                    "restart_counter": analysis.restart_counter}
+            return {'analysis_id': analysis.analysis_id,
+                    'project_id': analysis.project_id,
+                    'registry_url': analysis.registry_url,
+                    'image_url': analysis.image_url,
+                    'registry_user': analysis.registry_user,
+                    'registry_password': analysis.registry_password,
+                    'namespace': analysis.namespace,
+                    'kong_token': analysis.kong_token,
+                    'restart_counter': analysis.restart_counter}
         return None
 
     def delete_old_deployments_from_db(self, analysis_id: str) -> None:
         deployments = self.get_deployments(analysis_id)
         deployments = sorted(deployments, key=lambda x: x.time_created, reverse=True)
         for deployment in deployments[1:]:
-            self.delete_deployment(deployment.deployment_name)
+            self.delete_deployment(deployment.deployment_name)

src/resources/log/entity.py

@@ -10,7 +10,7 @@
 
 class LogEntity(BaseModel):
     log: str
-    log_type: Literal["emerg", "alert", "crit", "error", "warn", "notice", "info", "debug"]
+    log_type: Literal['emerg', 'alert', 'crit', 'error', 'warn', 'notice', 'info', 'debug']
 
     id: str = str(uuid.uuid4())
     created_at: str = str(datetime.now())
@@ -21,7 +21,7 @@ def __str__(self) -> str:
 
 class CreateLogEntity(BaseModel):
     log: str
-    log_type: Literal["emerg", "alert", "crit", "error", "warn", "notice", "info", "debug"]
+    log_type: Literal['emerg', 'alert', 'crit', 'error', 'warn', 'notice', 'info', 'debug']
 
     analysis_id: str
     status: str