add counter for internal paths calls and user agents

giograno · giograno · commit f78c5539751f · 2024-11-18T18:56:43.000-05:00
diff --git a/localstack-core/localstack/aws/handlers/validation.py b/localstack-core/localstack/aws/handlers/validation.py
@@ -3,6 +3,7 @@
 """
 
 import logging
+from typing import Any
 
 from openapi_core import OpenAPI
 from openapi_core.contrib.werkzeug import WerkzeugOpenAPIRequest, WerkzeugOpenAPIResponse
@@ -18,16 +19,43 @@
 from localstack.aws.chain import Handler, HandlerChain
 from localstack.constants import INTERNAL_RESOURCE_PATH
 from localstack.http import Response
+from localstack.utils.analytics.usage import UsageSetCounter
 
 LOG = logging.getLogger(__name__)
 
 
+class UsageCollectorFactory:
+    _collector_registry: dict[str, Any] = {}
+    """Registry for the different paths."""
+
+    NAMESPACE_PREFIX = "internal"
+    """Namespace prefix to track usage of public endpoints (_localstack/ and _aws/)."""
+
+    @classmethod
+    def get_collector(cls, path: str):
+        namespace = f"{cls.NAMESPACE_PREFIX}/{path}"
+        if namespace not in cls._collector_registry:
+            cls._collector_registry[namespace] = UsageSetCounter(namespace)
+        return cls._collector_registry[namespace]
+
+
 class OpenAPIValidator(Handler):
     open_apis: list["OpenAPI"]
 
     def __init__(self) -> None:
         self._load_specs()
 
+    def _record_usage(self, context: RequestContext) -> None:
+        if config.DISABLE_EVENTS:
+            return
+
+        request_path = context.request.path
+        user_agent = context.request.headers.get("User-Agent")
+        if not request_path or not user_agent:
+            return
+        collector = UsageCollectorFactory.get_collector(context.request.path)
+        collector.record(user_agent)
+
     def _load_specs(self) -> None:
         """Load the openapi spec plugins iff at least one between request and response validation is set."""
         if not (config.OPENAPI_VALIDATE_REQUEST or config.OPENAPI_VALIDATE_RESPONSE):
@@ -45,6 +73,7 @@ class OpenAPIRequestValidator(OpenAPIValidator):
     """
 
     def __call__(self, chain: HandlerChain, context: RequestContext, response: Response):
+        self._record_usage(context)
         if not config.OPENAPI_VALIDATE_REQUEST:
             return
 
