LITE-25146 add support for call context

Author: Francesco Faraone

connect/eaas/runner/handlers/anvil.py

@@ -8,11 +8,9 @@
 import anvil.server
 
 from connect.client import ConnectClient
-
-from pkg_resources import iter_entry_points
-
 from connect.eaas.core.logging import ExtensionLogHandler, RequestLogger
 from connect.eaas.runner.config import ConfigHelper
+from connect.eaas.runner.helpers import iter_entry_points
 
 
 logger = logging.getLogger(__name__)

connect/eaas/runner/handlers/events.py

@@ -1,12 +1,11 @@
 import inspect
 import logging
 
-from pkg_resources import iter_entry_points
-
 from connect.client import AsyncConnectClient, ConnectClient
 from connect.eaas.core.logging import ExtensionLogHandler, RequestLogger
 from connect.eaas.runner.config import ConfigHelper
 from connect.eaas.runner.constants import EVENT_TYPE_EXT_METHOD_MAP
+from connect.eaas.runner.helpers import iter_entry_points
 
 
 logger = logging.getLogger(__name__)

connect/eaas/runner/handlers/web.py

@@ -5,9 +5,9 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi.openapi.utils import get_openapi
-from pkg_resources import iter_entry_points
 
 from connect.eaas.runner.config import ConfigHelper
+from connect.eaas.runner.helpers import iter_entry_points
 
 
 logger = logging.getLogger(__name__)

connect/eaas/runner/helpers.py

@@ -8,15 +8,13 @@
 import subprocess
 import sys
 from datetime import datetime
+from importlib.metadata import entry_points, version
 from uuid import uuid4
 
 from connect.client import ClientError, ConnectClient
 
 import requests
-from pkg_resources import (
-    DistributionNotFound,
-    get_distribution,
-)
+
 
 from connect.eaas.runner.constants import (
     BACKGROUND_TASK_MAX_EXECUTION_TIME,
@@ -86,8 +84,8 @@ def get_environment():
 
 def get_version():
     try:
-        return get_distribution('connect-extension-runner').version
-    except DistributionNotFound:
+        return version('connect-extension-runner')
+    except Exception:
         return '0.0.0'
 
 
@@ -179,3 +177,15 @@ def notify_process_restarted(process_type):
             )
         except ClientError as ce:
             logger.warning(f'Cannot notify {process_type} process restart: {ce}')
+
+
+def iter_entry_points(group, name=None):
+    group_entrypoints = entry_points().get(group)
+    if not group_entrypoints:
+        return
+    for ep in group_entrypoints:
+        if name:
+            if ep.name == name:
+                yield ep
+        else:
+            yield ep

connect/eaas/runner/workers/web.py

@@ -72,6 +72,34 @@ async def shutdown(self):
     async def process_task(self, task):
         logger.info(f'new webtask received: {task.request.method} {task.request.url}')
         headers = copy.copy(task.request.headers)
+        headers.update(self.get_internal_headers(task))
+        message = None
+        try:
+            async with httpx.AsyncClient(
+                app=self.handler.app, base_url='http://localhost',
+            ) as client:
+                body = (
+                    base64.decodebytes(task.request.content.encode('utf-8'))
+                    if task.request.content else b''
+                )
+                response = await client.request(
+                    task.request.method,
+                    task.request.url,
+                    headers=headers,
+                    content=body,
+                )
+
+            message = self.build_response(
+                task, response.status_code, response.headers, response.content,
+            )
+        except Exception as e:
+            message = self.build_response(
+                task, 500, {}, str(e).encode('utf-8'),
+            )
+        await self.send(message)
+
+    def get_internal_headers(self, task):
+        headers = {}
         headers['X-Connect-Api-Gateway-Url'] = self.config.get_api_url()
         headers['X-Connect-User-Agent'] = self.config.get_user_agent()['User-Agent']
         headers['X-Connect-Extension-Id'] = self.config.service_id
@@ -83,26 +111,29 @@ async def process_task(self, task):
         if task.options.installation_id:
             headers['X-Connect-Installation-Id'] = task.options.installation_id
 
+        if task.options.connect_correlation_id:
+            headers['X-Connect-Correlation-Id'] = task.options.connect_correlation_id
+
+        if task.options.user_id:
+            headers['X-Connect-User-Id'] = task.options.user_id
+
+        if task.options.account_id:
+            headers['X-Connect-Account-Id'] = task.options.account_id
+
+        if task.options.account_role:
+            headers['X-Connect-Account-Role'] = task.options.account_role
+
+        if task.options.call_type:
+            headers['X-Connect-Call-Type'] = task.options.call_type
+
+        if task.options.call_source:
+            headers['X-Connect-Call-Source'] = task.options.call_source
+
         if self.config.logging_api_key is not None:
             headers['X-Connect-Logging-Api-Key'] = self.config.logging_api_key
             headers['X-Connect-Logging-Metadata'] = json.dumps(self.config.metadata)
 
-        async with httpx.AsyncClient(app=self.handler.app, base_url='http://localhost') as client:
-            body = (
-                base64.decodebytes(task.request.content.encode('utf-8'))
-                if task.request.content else b''
-            )
-            response = await client.request(
-                task.request.method,
-                task.request.url,
-                headers=headers,
-                content=body,
-            )
-
-        message = self.build_response(
-            task, response.status_code, response.headers, response.content,
-        )
-        await self.send(message)
+        return headers
 
     def build_response(self, task, status, headers, body):
         log = logger.info if status < 500 else logger.error