diff --git a/src/pydo/aio/operations/_patch.py b/src/pydo/aio/operations/_patch.py index 233de047..69c89a27 100644 --- a/src/pydo/aio/operations/_patch.py +++ b/src/pydo/aio/operations/_patch.py @@ -98,6 +98,26 @@ def __init__(self, *args, **kwargs): ) +# --------------------------------------------------------------------------- +# Insights PromQL form posts (async) +# --------------------------------------------------------------------------- + +try: + from ._operations import InsightsOperations as _GeneratedInsightsOperations + + _HAS_INSIGHTS = True +except ImportError: + _HAS_INSIGHTS = False + +if _HAS_INSIGHTS: + from pydo.custom_insights_prom import AsyncPromFormMixin + + class InsightsOperations( # pylint: disable=too-few-public-methods + AsyncPromFormMixin, _GeneratedInsightsOperations + ): + """Async Insights operations, including form-urlencoded PromQL posts.""" + + # --------------------------------------------------------------------------- # Exports # --------------------------------------------------------------------------- @@ -107,6 +127,8 @@ def __init__(self, *args, **kwargs): __all__.append("InferenceOperations") if _HAS_AGENT_INFERENCE: __all__.append("AgentInferenceOperations") +if _HAS_INSIGHTS: + __all__.append("InsightsOperations") def patch_sdk(): diff --git a/src/pydo/custom_insights_prom.py b/src/pydo/custom_insights_prom.py new file mode 100644 index 00000000..1a712d15 --- /dev/null +++ b/src/pydo/custom_insights_prom.py @@ -0,0 +1,185 @@ +"""Form-urlencoded PromQL posts. + +AutoRest's Python generator leaves ``application/x-www-form-urlencoded`` +operations unimplemented. These mixins supply the four Insights posts. +""" + +from typing import Any, Optional +from urllib.parse import quote + +from azure.core.exceptions import ( + ClientAuthenticationError, + HttpResponseError, + ResourceExistsError, + ResourceNotFoundError, + ResourceNotModifiedError, + map_error, +) +from azure.core.rest import HttpRequest + +_FORM = "application/x-www-form-urlencoded" +_RETURNED = {200, 400, 403, 404, 422, 500, 503} +_QUERY = "/v2/insights/query/{region}/prom/api/v1/query" +_QUERY_RANGE = "/v2/insights/query/{region}/prom/api/v1/query_range" +_LABELS = "/v2/insights/query/{region}/prom/api/v1/labels" +_SERIES = "/v2/insights/query/{region}/prom/api/v1/series" + + +def _form_request(path, region, fields, kwargs): + headers = dict(kwargs.pop("headers", {}) or {}) + params = kwargs.pop("params", None) or None + cls = kwargs.pop("cls", None) + error_map = { + 404: ResourceNotFoundError, + 409: ResourceExistsError, + 304: ResourceNotModifiedError, + 401: ClientAuthenticationError, + 429: HttpResponseError, + } + error_map.update(kwargs.pop("error_map", {}) or {}) + + url = path.format(region=quote(str(region), safe="")) + data = {key: value for key, value in fields.items() if value is not None} + if not any(key.lower() == "accept" for key in headers): + headers["Accept"] = "application/json" + if data: + request = HttpRequest("POST", url, params=params, headers=headers, data=data) + else: + headers["Content-Type"] = _FORM + request = HttpRequest("POST", url, params=params, headers=headers) + return request, cls, error_map + + +def _finish(ops, pipeline_response, cls, error_map): + response = pipeline_response.http_response + if response.status_code not in _RETURNED: + map_error( + status_code=response.status_code, + response=response, + error_map=error_map, + ) + raise HttpResponseError(response=response) + deserialized = response.json() if response.content else None + if cls: + headers = { + name: ops._deserialize( # pylint: disable=protected-access + "int", response.headers.get(name) + ) + for name in ( + "ratelimit-limit", + "ratelimit-remaining", + "ratelimit-reset", + ) + } + return cls(pipeline_response, deserialized, headers) + return deserialized + + +class PromFormMixin: + """Sync PromQL form posts. Subclass the generated Insights operations.""" + + def _prom_post(self, path, region, fields, **kwargs): + request, cls, error_map = _form_request(path, region, fields, kwargs) + request.url = self._client.format_url( + request.url + ) # pylint: disable=protected-access + pipeline_response = ( + self._client._pipeline.run( # pylint: disable=protected-access + request, stream=False, **kwargs + ) + ) + return _finish(self, pipeline_response, cls, error_map) + + def post_prom_query( + self, + region: str, + *, + query: str, + time: Optional[str] = None, + timeout: Optional[str] = None, + **kwargs: Any, + ) -> Any: + """POST an instant PromQL query as ``application/x-www-form-urlencoded``.""" + return self._prom_post( + _QUERY, + region, + {"query": query, "time": time, "timeout": timeout}, + **kwargs, + ) + + def post_prom_query_range( + self, + region: str, + *, + query: str, + start: str, + end: str, + step: str, + timeout: Optional[str] = None, + **kwargs: Any, + ) -> Any: + """POST a range PromQL query as ``application/x-www-form-urlencoded``.""" + return self._prom_post( + _QUERY_RANGE, + region, + { + "query": query, + "start": start, + "end": end, + "step": step, + "timeout": timeout, + }, + **kwargs, + ) + + def post_prom_labels( + self, + region: str, + *, + start: Optional[str] = None, + end: Optional[str] = None, + match: Optional[list] = None, + **kwargs: Any, + ) -> Any: + """POST label-name selectors as ``application/x-www-form-urlencoded``.""" + return self._prom_post( + _LABELS, + region, + {"start": start, "end": end, "match[]": match}, + **kwargs, + ) + + def post_prom_series( + self, + region: str, + *, + match: list, + start: Optional[str] = None, + end: Optional[str] = None, + **kwargs: Any, + ) -> Any: + """POST series selectors as ``application/x-www-form-urlencoded``.""" + return self._prom_post( + _SERIES, + region, + {"match[]": match, "start": start, "end": end}, + **kwargs, + ) + + +class AsyncPromFormMixin(PromFormMixin): + """Async PromQL form posts. Public methods return awaitables.""" + + async def _prom_post( # pylint: disable=invalid-overridden-method + self, path, region, fields, **kwargs + ): + request, cls, error_map = _form_request(path, region, fields, kwargs) + request.url = self._client.format_url( + request.url + ) # pylint: disable=protected-access + pipeline_response = ( + await self._client._pipeline.run( # pylint: disable=protected-access + request, stream=False, **kwargs + ) + ) + return _finish(self, pipeline_response, cls, error_map) diff --git a/src/pydo/operations/_patch.py b/src/pydo/operations/_patch.py index 9a3ef3ce..84c86d81 100644 --- a/src/pydo/operations/_patch.py +++ b/src/pydo/operations/_patch.py @@ -111,6 +111,28 @@ def __init__(self, *args, **kwargs): ) +# --------------------------------------------------------------------------- +# Insights PromQL form posts +# --------------------------------------------------------------------------- +# AutoRest marks application/x-www-form-urlencoded operations abstract. +# This subclass exists once `make generate` emits InsightsOperations. + +try: + from ._operations import InsightsOperations as _GeneratedInsightsOperations + + _HAS_INSIGHTS = True +except ImportError: + _HAS_INSIGHTS = False + +if _HAS_INSIGHTS: + from pydo.custom_insights_prom import PromFormMixin + + class InsightsOperations( # pylint: disable=too-few-public-methods + PromFormMixin, _GeneratedInsightsOperations + ): + """Insights operations, including form-urlencoded PromQL posts.""" + + # --------------------------------------------------------------------------- # Exports # --------------------------------------------------------------------------- @@ -120,6 +142,8 @@ def __init__(self, *args, **kwargs): __all__.append("InferenceOperations") if _HAS_AGENT_INFERENCE: __all__.append("AgentInferenceOperations") +if _HAS_INSIGHTS: + __all__.append("InsightsOperations") def patch_sdk(): diff --git a/tests/mocked/test_insights_prom_form.py b/tests/mocked/test_insights_prom_form.py new file mode 100644 index 00000000..333d6607 --- /dev/null +++ b/tests/mocked/test_insights_prom_form.py @@ -0,0 +1,192 @@ +"""Handwritten PromQL form posts that AutoRest leaves unimplemented.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring +# pylint: disable=too-few-public-methods,protected-access,unused-argument +# pylint: disable=invalid-overridden-method + +import asyncio + +import pytest +from azure.core.exceptions import ClientAuthenticationError, HttpResponseError + +from pydo.custom_insights_prom import AsyncPromFormMixin, PromFormMixin + + +def _not_implemented(cls, methods): + missing = [name for name in methods if not callable(getattr(cls, name, None))] + if missing: + raise NotImplementedError(missing) + + +class _Pipeline: + def __init__(self, status_code=200, body=None): + self.request = None + self.kwargs = None + self.status_code = status_code + self.body = {"status": "success"} if body is None else body + + def run(self, request, **kwargs): + self.request = request + self.kwargs = kwargs + return _PipelineResponse(self.status_code, self.body) + + +class _PipelineResponse: + def __init__(self, status_code, body): + self.http_response = _Response(status_code, body) + + +class _Response: + def __init__(self, status_code, body): + self.status_code = status_code + self.reason = "error" + self._body = body + self.headers = {} + self.content = b"{}" if body is not None else b"" + + def json(self): + return self._body + + +class _Client: + def __init__(self, pipeline): + self._pipeline = pipeline + + def format_url(self, url): + return "https://api.digitalocean.com" + url + + +def _ops(status_code=200, body=None): + pipeline = _Pipeline(status_code, body) + return PromFormMixin(), pipeline, _Client(pipeline) + + +def _bind(ops, client): + ops._client = client # pylint: disable=protected-access + ops._deserialize = ( + lambda *_args, **_kwargs: None + ) # pylint: disable=protected-access + return ops + + +def test_post_prom_query_sends_form_body(): + ops, pipeline, client = _ops() + _bind(ops, client) + + body = ops.post_prom_query("nyc3", query="up", time="1620683817") + + request = pipeline.request + assert body == {"status": "success"} + assert request.method == "POST" + assert ( + request.url + == "https://api.digitalocean.com/v2/insights/query/nyc3/prom/api/v1/query" + ) + assert request.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert request.content == {"query": "up", "time": "1620683817"} + assert pipeline.kwargs["stream"] is False + + +def test_post_prom_query_range_sends_required_fields(): + ops, pipeline, client = _ops() + _bind(ops, client) + + ops.post_prom_query_range( + "nyc3", query="up", start="1", end="2", step="15s", timeout="30s" + ) + + assert pipeline.request.url.endswith("/prom/api/v1/query_range") + assert pipeline.request.content == { + "query": "up", + "start": "1", + "end": "2", + "step": "15s", + "timeout": "30s", + } + + +def test_post_prom_labels_omits_missing_fields(): + ops, pipeline, client = _ops() + _bind(ops, client) + + ops.post_prom_labels("nyc3", match=["up"]) + + assert pipeline.request.url.endswith("/prom/api/v1/labels") + assert pipeline.request.content == {"match[]": ["up"]} + + +def test_post_prom_series_repeats_match_selectors(): + ops, pipeline, client = _ops() + _bind(ops, client) + + ops.post_prom_series("nyc3", match=["up", "go_goroutines"], start="1") + + assert pipeline.request.url.endswith("/prom/api/v1/series") + assert pipeline.request.content == { + "match[]": ["up", "go_goroutines"], + "start": "1", + } + + +def test_prometheus_error_body_is_returned(): + ops, _pipeline, client = _ops(400, {"status": "error", "errorType": "bad_data"}) + _bind(ops, client) + + body = ops.post_prom_query("nyc3", query="???") + + assert body["errorType"] == "bad_data" + + +def test_unauthorized_raises(): + ops, _pipeline, client = _ops(401, {"id": "unauthorized"}) + _bind(ops, client) + + with pytest.raises(ClientAuthenticationError): + ops.post_prom_query("nyc3", query="up") + + +def test_unexpected_status_raises(): + ops, _pipeline, client = _ops(418, {}) + _bind(ops, client) + + with pytest.raises(HttpResponseError): + ops.post_prom_query("nyc3", query="up") + + +def test_subclass_satisfies_autorest_unimplemented_check(): + class _Generated: + def __init__(self, *args, **kwargs): + _not_implemented( + self.__class__, + [ + "post_prom_query", + "post_prom_query_range", + "post_prom_labels", + "post_prom_series", + ], + ) + + class InsightsOperations(PromFormMixin, _Generated): + pass + + InsightsOperations(object(), object(), object(), object()) + + +def test_async_post_prom_query_awaits_pipeline(): + class _AsyncPipeline(_Pipeline): + async def run(self, request, **kwargs): + self.request = request + self.kwargs = kwargs + return _PipelineResponse(self.status_code, self.body) + + pipeline = _AsyncPipeline() + ops = _bind(AsyncPromFormMixin(), _Client(pipeline)) + + body = asyncio.get_event_loop().run_until_complete( + ops.post_prom_query("nyc3", query="up") + ) + + assert body == {"status": "success"} + assert ( + pipeline.request.headers["Content-Type"] == "application/x-www-form-urlencoded" + )