Merge pull request #436 from xchem/model-refactoring-stage-1

Collects changes from underlying branch

viewer/services.py

@@ -0,0 +1,182 @@
+import os
+import logging
+from enum import Enum
+
+import asyncio
+from concurrent import futures
+import functools
+import requests
+
+from pydiscourse import DiscourseClient
+
+from viewer.squonk2_agent import get_squonk2_agent
+from frag.utils.network_utils import get_driver
+from api import security
+
+logger = logging.getLogger(__name__)
+
+DELAY = 10
+
+
+class State(str, Enum):
+    NOT_CONFIGURED = "NOT_CONFIGURED"
+    DEGRADED = "DEGRADED"
+    OK = "OK"
+    TIMEOUT = "TIMEOUT"
+    ERROR = "ERROR"
+
+
+class Service(str, Enum):
+    ISPYB = "ispyb"
+    DISCOURSE = "discourse"
+    SQUONK = "squonk"
+    FRAG = "fragmentation_graph"
+    KEYCLOAK = "keycloak"
+
+
+# called from the outside
+def get_service_state(services):
+    return asyncio.run(service_queries(services))
+
+
+async def service_queries(services):
+    logger.debug("service query called")
+    coroutines = []
+    if Service.ISPYB in services:
+        coroutines.append(
+            ispyb(
+                Service.ISPYB,
+                "Access control (ISPyB)",
+                ispyb_host="ISPYB_HOST",
+            )
+        )
+
+    if Service.SQUONK in services:
+        coroutines.append(
+            squonk(
+                Service.SQUONK,
+                "Squonk",
+                squonk_pwd="SQUONK2_ORG_OWNER_PASSWORD",
+            )
+        )
+
+    if Service.FRAG in services:
+        coroutines.append(
+            fragmentation_graph(
+                Service.FRAG,
+                "Fragmentation graph",
+                url="NEO4J_BOLT_URL",
+            )
+        )
+
+    if Service.DISCOURSE in services:
+        coroutines.append(
+            discourse(
+                Service.DISCOURSE,
+                "Discourse",
+                key="DISCOURSE_API_KEY",
+                url="DISCOURSE_HOST",
+                user="DISCOURSE_USER",
+            )
+        )
+
+    if Service.KEYCLOAK in services:
+        coroutines.append(
+            keycloak(
+                Service.KEYCLOAK,
+                "Keycloak",
+                url="OIDC_KEYCLOAK_REALM",
+                secret="OIDC_RP_CLIENT_SECRET",
+            )
+        )
+
+    result = await asyncio.gather(*coroutines)
+    logger.debug("service query result: %s", result)
+    return result
+
+
+def service_query(func):
+    @functools.wraps(func)
+    async def wrapper_service_query(*args, **kwargs):
+        logger.debug("+ wrapper_service_query")
+        logger.debug("args passed: %s", args)
+        logger.debug("kwargs passed: %s", kwargs)
+
+        state = State.NOT_CONFIGURED
+        envs = [os.environ.get(k, None) for k in kwargs.values()]
+        # env variables come in kwargs, all must be defined
+        if all(envs):
+            state = State.DEGRADED
+            loop = asyncio.get_running_loop()
+            executor = futures.ThreadPoolExecutor()
+            try:
+                async with asyncio.timeout(DELAY):
+                    future = loop.run_in_executor(
+                        executor, functools.partial(func, *args, **kwargs)
+                    )
+                    conn = await future
+                    if conn:
+                        state = State.OK
+
+            except TimeoutError:
+                logger.error("Service query function %s timed out", func.__name__)
+                state = State.TIMEOUT
+            except Exception as exc:
+                # unknown error with the query
+                logger.exception(exc, exc_info=True)
+                state = State.ERROR
+
+        # name and ID are 2nd and 1st params respectively (0 is
+        # self). alternative solution for this would be to return
+        # just a state and have the service_queries() map the
+        # results to the correct values
+        return {"id": args[0], "name": args[1], "state": state}
+
+    return wrapper_service_query
+
+
+@service_query
+async def ispyb(func_id, name, ispyb_host=None):
+    logger.debug("+ ispyb")
+    return security.get_conn()
+
+
+@service_query
+async def discourse(func_id, name, key=None, url=None, user=None):
+    logger.debug("+ discourse")
+    client = DiscourseClient(
+        os.environ.get(url, None),
+        api_username=os.environ.get(user, None),
+        api_key=os.environ.get(key, None),
+    )
+    # TODO: some action on client?
+    return client
+
+
+@service_query
+async def squonk(func_id, name, squonk_pwd=None):
+    logger.debug("+ squonk")
+    return get_squonk2_agent().configured().success
+
+
+@service_query
+async def fragmentation_graph(func_id, name, url=None):
+    logger.debug("+ fragmentation_graph")
+    # graph_driver = get_driver(url='neo4j', neo4j_auth='neo4j/password')
+    graph_driver = get_driver()
+    with graph_driver.session() as session:
+        try:
+            _ = session.run("match (n) return count (n);")
+            return True
+        except ValueError:
+            # service isn't running
+            return False
+
+
+@service_query
+async def keycloak(func_id, name, url=None, secret=None):
+    logger.debug("+ keycloak")
+    keycloak_realm = os.environ.get(url, None)
+    response = requests.get(keycloak_realm)
+    logger.debug("keycloak response: %s", response)
+    return response.ok

viewer/views.py

@@ -12,9 +12,6 @@
 import pytz
 from pathlib import Path
 
-from enum import Enum
-
-from pydiscourse import DiscourseClient
 
 import pandas as pd
 
@@ -39,11 +36,8 @@
 
 from celery.result import AsyncResult
 
-from api import security
 from api.security import ISpyBSafeQuerySet
 
-from frag.utils.network_utils import get_driver
-
 from api.utils import get_params, get_highlighted_diffs, pretty_request
 from viewer.utils import create_squonk_job_request_url
 from viewer.utils import handle_uploaded_file
@@ -54,6 +48,7 @@
 from viewer.squonk2_agent import Squonk2AgentRv, Squonk2Agent, get_squonk2_agent
 from viewer.squonk2_agent import AccessParams, CommonParams, SendParams, RunJobParams
 
+from viewer.services import get_service_state
 
 
 from .forms import CSetForm
@@ -2175,11 +2170,6 @@ def get(self, request):
 
 class ServiceState(View):
 
-    class State(str, Enum):
-        NOT_CONFIGURED = "NOT_CONFIGURED"
-        DEGRADED = "DEGRADED"
-        OK = "OK"
-
     def get(self, *args, **kwargs):
         """Poll external service status.
 
@@ -2199,78 +2189,6 @@ def get(self, *args, **kwargs):
         services = [k for k in service_string.split(":") if k !=  ""]
         logger.debug("Services ordered: %s", services)
 
-        service_states = []
-        for s in services:
-            func = getattr(self, s, None)
-            if func:
-                service_states.append(func())
+        service_states = get_service_state(services)
 
         return JsonResponse({"service_states": service_states})
-
-
-    @staticmethod
-    def ispyb():
-        state = ServiceState.State.NOT_CONFIGURED
-        if os.environ.get("ISPYB_HOST", None):
-            state = ServiceState.State.DEGRADED
-            if security.get_conn():
-                state = ServiceState.State.OK
-
-        return {"id": "ispyb", "name": "Access control (ISPyB)", "state": state}
-
-
-    @staticmethod
-    def squonk():
-        state = ServiceState.State.NOT_CONFIGURED
-        if os.environ.get("SQUONK2_ORG_OWNER_PASSWORD", None):
-            state = ServiceState.State.DEGRADED
-            if get_squonk2_agent().configured().success:
-                state = ServiceState.State.OK
-
-        return {"id": "squonk", "name": "Squonk", "state": state}
-
-    @staticmethod
-    def fragmentation_graph():
-        state = ServiceState.State.NOT_CONFIGURED
-        if os.environ.get("NEO4J_BOLT_URL", None):
-            state = ServiceState.State.DEGRADED
-            # graph_driver = get_driver(url='neo4j', neo4j_auth='neo4j/password')
-            graph_driver = get_driver()
-            with graph_driver.session() as session:
-                try:
-                    _ = session.run('match (n) return count (n);')
-                    state = ServiceState.State.OK
-                except ValueError:
-                    # service isn't running
-                    pass
-
-        return {"id": "fragmentation_graph", "name": "Fragmentation graph", "state": state}
-
-    @staticmethod
-    def discourse():
-        state = ServiceState.State.NOT_CONFIGURED
-        key = os.environ.get("DISCOURSE_API_KEY", None)
-        url = os.environ.get("DISCOURSE_HOST", None)
-        user = os.environ.get("DISCOURSE_USER", None)
-        if key and url and user:
-            state = ServiceState.State.DEGRADED
-            client = DiscourseClient(
-                url,
-                api_username=user,
-                api_key=key,
-            )
-            # TODO: some action on client?
-            if client:
-                state = ServiceState.State.OK
-
-        return {"id": "discourse", "name": "Discourse", "state": state}
-
-    @staticmethod
-    def keycloak():
-        state = ServiceState.State.NOT_CONFIGURED
-        if os.environ.get("OIDC_RP_CLIENT_SECRET", None):
-            state = ServiceState.State.DEGRADED
-            if security.get_conn():
-                state = ServiceState.State.OK
-
-        return {"id": "keycloak", "name": "Keycloak", "state": state}