feat: gunicorn (#591)

shahargl · web-flow · commit 1c3d63b7a08e · 2023-11-30T14:46:52.000+02:00
diff --git a/docker/Dockerfile.api b/docker/Dockerfile.api
@@ -22,4 +22,4 @@ FROM base as final
 ENV PATH="/venv/bin:${PATH}"
 ENV VIRTUAL_ENV="/venv"
 COPY --from=builder /venv /venv
-ENTRYPOINT ["keep", "--json", "api"]
+ENTRYPOINT ["gunicorn", "keep.api.api:get_app", "--bind" , "0.0.0.0:8080" , "--workers", "4" , "-k" , "uvicorn.workers.UvicornWorker" ]
diff --git a/keep-ui/app/settings/smtp-settings.tsx b/keep-ui/app/settings/smtp-settings.tsx
@@ -3,6 +3,7 @@ import { Card, Button, Title, Subtitle, TextInput } from "@tremor/react";
 import useSWR from "swr";
 import { getApiURL } from "utils/apiUrl";
 import { fetcher } from "utils/fetcher";
+import Loading from "app/loading";
 
 interface SMTPSettings {
   host: string;
@@ -91,8 +92,9 @@ export default function SMTPSettingsForm({ accessToken, selectedTab }: Props) {
   );
 
   // Show loading state or error messages if needed
-  if (isLoading) return <div>Loading...</div>; // Loading state
-  if (error) return <div>Error: {error.message}</div>;
+  if (smtpSettings === undefined || isLoading) {
+    return <Loading />;
+  }
 
   // if no errors and we have data, set the settings
   if (smtpSettings) {
diff --git a/keep/api/api.py b/keep/api/api.py
@@ -9,6 +9,7 @@
 from dotenv import find_dotenv, load_dotenv
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.responses import JSONResponse
+from opentelemetry import trace
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.middleware.cors import CORSMiddleware
 from starlette_context import plugins
@@ -61,6 +62,7 @@ class EventCaptureMiddleware(BaseHTTPMiddleware):
     def __init__(self, app: FastAPI):
         super().__init__(app)
         self.posthog_client = get_posthog_client()
+        self.tracer = trace.get_tracer(__name__)
 
     def _extract_identity(self, request: Request) -> str:
         try:
@@ -70,42 +72,45 @@ def _extract_identity(self, request: Request) -> str:
         except Exception:
             return "anonymous"
 
-    def capture_request(self, request: Request) -> None:
+    async def capture_request(self, request: Request) -> None:
         identity = self._extract_identity(request)
-        self.posthog_client.capture(
-            identity,
-            "request-started",
-            {"path": request.url.path, "method": request.method},
-        )
+        with self.tracer.start_as_current_span("capture_request"):
+            self.posthog_client.capture(
+                identity,
+                "request-started",
+                {"path": request.url.path, "method": request.method},
+            )
 
-    def capture_response(self, request: Request, response: Response) -> None:
+    async def capture_response(self, request: Request, response: Response) -> None:
         identity = self._extract_identity(request)
-        self.posthog_client.capture(
-            identity,
-            "request-finished",
-            {
-                "path": request.url.path,
-                "method": request.method,
-                "status_code": response.status_code,
-            },
-        )
+        with self.tracer.start_as_current_span("capture_response"):
+            self.posthog_client.capture(
+                identity,
+                "request-finished",
+                {
+                    "path": request.url.path,
+                    "method": request.method,
+                    "status_code": response.status_code,
+                },
+            )
 
-    def flush(self):
-        logger.info("Flushing Posthog events")
-        self.posthog_client.flush()
-        logger.info("Posthog events flushed")
+    async def flush(self):
+        with self.tracer.start_as_current_span("flush_posthog_events"):
+            logger.info("Flushing Posthog events")
+            self.posthog_client.flush()
+            logger.info("Posthog events flushed")
 
     async def dispatch(self, request: Request, call_next):
         # Skip OPTIONS requests
         if request.method == "OPTIONS":
             return await call_next(request)
         # Capture event before request
-        self.capture_request(request)
+        await self.capture_request(request)
 
         response = await call_next(request)
 
         # Capture event after request
-        self.capture_response(request, response)
+        await self.capture_response(request, response)
 
         # Perform async tasks or flush events after the request is handled
         self.flush()
@@ -219,6 +224,13 @@ async def on_startup():
             ] = verify_token_or_key_single_tenant
             try_create_single_tenant(SINGLE_TENANT_UUID)
 
+        # load all providers into cache
+        from keep.providers.providers_factory import ProvidersFactory
+
+        logger.info("Loading providers into cache")
+        ProvidersFactory.get_all_providers()
+        logger.info("Providers loaded successfully")
+
     @app.exception_handler(Exception)
     async def catch_exception(request: Request, exc: Exception):
         logging.error(
@@ -233,6 +245,13 @@ async def catch_exception(request: Request, exc: Exception):
             },
         )
 
+    @app.middleware("http")
+    async def log_middeware(request: Request, call_next):
+        logger.info(f"Request started: {request.method} {request.url.path}")
+        response = await call_next(request)
+        logger.info(f"Request finished: {request.method} {request.url.path}")
+        return response
+
     keep.api.observability.setup(app)
 
     if os.environ.get("USE_NGROK", "false") == "true":
diff --git a/keep/api/core/db.py b/keep/api/core/db.py
@@ -119,8 +119,12 @@ def get_session() -> Session:
     Yields:
         Session: A database session
     """
-    with Session(engine) as session:
-        yield session
+    from opentelemetry import trace
+
+    tracer = trace.get_tracer(__name__)
+    with tracer.start_as_current_span("get_session"):
+        with Session(engine) as session:
+            yield session
 
 
 def try_create_single_tenant(tenant_id: str) -> None:
diff --git a/keep/api/core/dependencies.py b/keep/api/core/dependencies.py
@@ -33,18 +33,22 @@
 
 
 def get_user_email(request: Request) -> str | None:
-    token = request.headers.get("Authorization")
-    if token:
-        token = token.split(" ")[1]
-        decoded_token = jwt.decode(token, options={"verify_signature": False})
-        return decoded_token.get("email")
-    elif "x-api-key" in request.headers:
-        username = get_user_by_api_key(request.headers["x-api-key"])
-        return username
-    else:
-        raise HTTPException(
-            status_code=401, detail="Invalid authentication credentials"
-        )
+    from opentelemetry import trace
+
+    tracer = trace.get_tracer(__name__)
+    with tracer.start_as_current_span("get_user_email"):
+        token = request.headers.get("Authorization")
+        if token:
+            token = token.split(" ")[1]
+            decoded_token = jwt.decode(token, options={"verify_signature": False})
+            return decoded_token.get("email")
+        elif "x-api-key" in request.headers:
+            username = get_user_by_api_key(request.headers["x-api-key"])
+            return username
+        else:
+            raise HTTPException(
+                status_code=401, detail="Invalid authentication credentials"
+            )
 
 
 def __extract_api_key(
@@ -133,34 +137,42 @@ def verify_api_key(
     return tenant_api_key.tenant_id
 
 
+# init once so the cache will actually work
+auth_domain = os.environ.get("AUTH0_DOMAIN")
+if auth_domain:
+    jwks_uri = f"https://{auth_domain}/.well-known/jwks.json"
+    jwks_client = jwt.PyJWKClient(jwks_uri, cache_keys=True)
+
+
 def verify_bearer_token(token: str = Depends(oauth2_scheme)) -> str:
     # Took the implementation from here:
     #   https://github.com/auth0-developer-hub/api_fastapi_python_hello-world/blob/main/application/json_web_token.py
-    if not token:
-        raise HTTPException(status_code=401, detail="No token provided 👈")
-    try:
-        auth_domain = os.environ.get("AUTH0_DOMAIN")
-        auth_audience = os.environ.get("AUTH0_AUDIENCE")
-        jwks_uri = f"https://{auth_domain}/.well-known/jwks.json"
-        issuer = f"https://{auth_domain}/"
-        jwks_client = jwt.PyJWKClient(jwks_uri)
-        jwt_signing_key = jwks_client.get_signing_key_from_jwt(token).key
-        payload = jwt.decode(
-            token,
-            jwt_signing_key,
-            algorithms="RS256",
-            audience=auth_audience,
-            issuer=issuer,
-            leeway=60,
-        )
-        tenant_id = payload.get("keep_tenant_id")
-        return tenant_id
-    except jwt.exceptions.DecodeError:
-        logger.exception("Failed to decode token")
-        raise HTTPException(status_code=401, detail="Token is not a valid JWT")
-    except Exception as e:
-        logger.exception("Failed to validate token")
-        raise HTTPException(status_code=401, detail=str(e))
+    from opentelemetry import trace
+
+    tracer = trace.get_tracer(__name__)
+    with tracer.start_as_current_span("verify_bearer_token"):
+        if not token:
+            raise HTTPException(status_code=401, detail="No token provided 👈")
+        try:
+            auth_audience = os.environ.get("AUTH0_AUDIENCE")
+            issuer = f"https://{auth_domain}/"
+            jwt_signing_key = jwks_client.get_signing_key_from_jwt(token).key
+            payload = jwt.decode(
+                token,
+                jwt_signing_key,
+                algorithms="RS256",
+                audience=auth_audience,
+                issuer=issuer,
+                leeway=60,
+            )
+            tenant_id = payload.get("keep_tenant_id")
+            return tenant_id
+        except jwt.exceptions.DecodeError:
+            logger.exception("Failed to decode token")
+            raise HTTPException(status_code=401, detail="Token is not a valid JWT")
+        except Exception as e:
+            logger.exception("Failed to validate token")
+            raise HTTPException(status_code=401, detail=str(e))
 
 
 def get_user_email_single_tenant(request: Request) -> str:
diff --git a/keep/api/logging.py b/keep/api/logging.py
@@ -67,7 +67,7 @@ def dump(self):
     "disable_existing_loggers": False,
     "formatters": {
         "json": {
-            "format": "%(asctime)s %(message)s %(levelname)s %(name)s %(filename)s %(otelTraceID)s %(otelSpanID)s %(otelServiceName)s",
+            "format": "%(asctime)s %(message)s %(levelname)s %(name)s %(filename)s %(otelTraceID)s %(otelSpanID)s %(otelServiceName)s %(threadName)s %(process)s",
             "class": "pythonjsonlogger.jsonlogger.JsonFormatter",
         }
     },
diff --git a/keep/api/routes/settings.py b/keep/api/routes/settings.py
@@ -1,5 +1,6 @@
 import io
 import json
+import logging
 import os
 import secrets
 import smtplib
@@ -27,6 +28,8 @@
 
 router = APIRouter()
 
+logger = logging.getLogger(__name__)
+
 
 @router.get(
     "/webhook",
@@ -36,6 +39,7 @@ def webhook_settings(
     tenant_id: str = Depends(verify_bearer_token),
     session: Session = Depends(get_session),
 ) -> WebhookSettings:
+    logger.info("Getting webhook settings")
     api_url = config("KEEP_API_URL")
     keep_webhook_api_url = f"{api_url}/alerts/event"
     webhook_api_key = get_or_create_api_key(
@@ -45,6 +49,7 @@ def webhook_settings(
         unique_api_key_id="webhook",
         system_description="Webhooks API key",
     )
+    logger.info("Webhook settings retrieved successfully")
     return WebhookSettings(
         webhookApi=keep_webhook_api_url,
         apiKey=webhook_api_key,
@@ -180,12 +185,14 @@ async def get_smtp_settings(
     tenant_id: str = Depends(verify_bearer_token),
     session: Session = Depends(get_session),
 ):
+    logger.info("Getting SMTP settings")
     context_manager = ContextManager(tenant_id=tenant_id)
     secret_manager = SecretManagerFactory.get_secret_manager(context_manager)
     # Read the SMTP settings from the secret manager
     try:
         smtp_settings = secret_manager.read_secret(secret_name="smtp")
         smtp_settings = json.loads(smtp_settings)
+        logger.info("SMTP settings retrieved successfully")
         return JSONResponse(status_code=200, content=smtp_settings)
     except Exception:
         # everything ok but no smtp settings
@@ -278,6 +285,7 @@ def get_api_key(
     session: Session = Depends(get_session),
     user_name: str = Depends(get_user_email),
 ):
+    logger.info("Getting API key")
     # get the api key for the CLI
     api_key = get_or_create_api_key(
         session=session,
@@ -286,4 +294,5 @@ def get_api_key(
         unique_api_key_id="cli",
         system_description="API key",
     )
+    logger.info("API key retrieved successfully")
     return {"apiKey": api_key}
diff --git a/keep/contextmanager/contextmanager.py b/keep/contextmanager/contextmanager.py
@@ -15,7 +15,7 @@ class ContextManager:
     STATE_FILE = "keepstate.json"
 
     def __init__(
-        self, tenant_id, workflow_id=None, workflow_execution_id=None, load_state=True
+        self, tenant_id, workflow_id=None, workflow_execution_id=None, load_state=False
     ):
         self.logger = logging.getLogger(__name__)
         self.logger_adapter = WorkflowLoggerAdapter(
@@ -37,7 +37,7 @@ def __init__(
         except RuntimeError:
             self.click_context = {}
         self.aliases = {}
-        self.state = {}
+        self._state = {}
         # dependencies are used so iohandler will be able to use the output class of the providers
         # e.g. let's say bigquery_provider results are google.cloud.bigquery.Row
         #     and we want to use it in iohandler, we need to import it before the eval
@@ -190,7 +190,7 @@ def set_step_context(self, step_id, results, foreach=False):
 
     def __load_state(self):
         try:
-            self.state = json.loads(
+            self._state = json.loads(
                 self.storage_manager.get_file(
                     self.tenant_id, self.state_file, create_if_not_exist=True
                 )
@@ -201,7 +201,13 @@ def __load_state(self):
                 f"State storage: {self.storage_manager.__class__.__name__}"
             )
             self.logger.warning(f"Reason: {exc}")
-            self.state = {}
+            self._state = {}
+
+    @property
+    def state(self):
+        if not self._state:
+            self.__load_state()
+        return self._state
 
     def get_last_workflow_run(self, workflow_id):
         if workflow_id in self.state:
diff --git a/keep/providers/providers_factory.py b/keep/providers/providers_factory.py
@@ -24,6 +24,8 @@ class ProviderConfigurationException(Exception):
 
 
 class ProvidersFactory:
+    _loaded_providers_cache = None
+
     @staticmethod
     def get_provider_class(provider_type: str) -> BaseProvider:
         provider_type_split = provider_type.split(
@@ -157,6 +159,13 @@ def get_all_providers() -> list[Provider]:
         Returns:
             list: All the providers.
         """
+        logger = logging.getLogger(__name__)
+        # use the cache if exists
+        if ProvidersFactory._loaded_providers_cache:
+            logger.info("Using cached providers")
+            return ProvidersFactory._loaded_providers_cache
+
+        logger.info("Loading providers")
         providers = []
         blacklisted_providers = [
             "base_provider",
@@ -278,6 +287,8 @@ def get_all_providers() -> list[Provider]:
             except ModuleNotFoundError:
                 logger.exception(f"Cannot import provider {provider_directory}")
                 continue
+
+        ProvidersFactory._loaded_providers_cache = providers
         return providers
 
     @staticmethod
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml

Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,7 @@ def dump(self):`
`67`	`67`	`"disable_existing_loggers": False,`
`68`	`68`	`"formatters": {`
`69`	`69`	`"json": {`
`70`		`- "format": "%(asctime)s %(message)s %(levelname)s %(name)s %(filename)s %(otelTraceID)s %(otelSpanID)s %(otelServiceName)s",`
	`70`	`+ "format": "%(asctime)s %(message)s %(levelname)s %(name)s %(filename)s %(otelTraceID)s %(otelSpanID)s %(otelServiceName)s %(threadName)s %(process)s",`
`71`	`71`	`"class": "pythonjsonlogger.jsonlogger.JsonFormatter",`
`72`	`72`	`}`
`73`	`73`	`},`