File size: 5,592 Bytes
79879d4
 
 
 
 
 
b7dddbe
79879d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f8c4791
79879d4
 
 
 
 
f8c4791
 
 
 
79879d4
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
from __future__ import annotations

import asyncio
import time
from typing import Any, Dict, List, Optional

from fastapi import APIRouter
from pydantic import BaseModel, Field, field_validator

from app.core.logger import get_logger
from app.core.thread_pool import thread_pool
from app.services.keys_extractor_service import KeysExtractor

logger = get_logger(__name__)

router = APIRouter()


class KeysExtractRequest(BaseModel):
    data: Any = Field(..., description="JSON object or array to search")
    key_names: Optional[List[str]] = Field(
        default=None,
        description=(
            "Key names to look up at any depth. Mutually exclusive with `query`; "
            "provide one of the two."
        ),
    )
    query: Optional[str] = Field(
        default=None,
        description=(
            "Extended JSONPath expression (parsed with `jsonpath_ng.ext`) used to "
            "select values. Supports filters, arithmetic and the union operator, e.g. "
            "`$.store.book[?price > 10].title`. Each match is returned with its "
            "full path inside the source document. Mutually exclusive with `key_names`."
        ),
    )
    result_limit: Optional[int] = Field(
        None,
        ge=1,
        description=(
            "Maximum results to return. For `key_names` mode this caps values per "
            "key; for `query` mode it caps the number of matches returned. "
            "Omit (or pass null) to return the full, uncapped result."
        ),
    )

    @field_validator("key_names")
    @classmethod
    def _validate_key_names(cls, v: Optional[List[str]]) -> Optional[List[str]]:
        if v is None:
            return v
        if not v:
            raise ValueError("key_names must be a non-empty list (omit it or supply a `query`)")
        for kn in v:
            if not isinstance(kn, str) or not kn:
                raise ValueError("each key_name must be a non-empty string")
        return v

    @field_validator("query")
    @classmethod
    def _validate_query(cls, v: Optional[str]) -> Optional[str]:
        if v is None:
            return v
        if not isinstance(v, str) or not v.strip():
            raise ValueError("query must be a non-empty string")
        return v


class KeysExtractResponse(BaseModel):
    success: bool
    time_ms: float
    data: Dict[str, Any]
    error_message: Optional[str] = None


class KeysExtractBatchRequest(BaseModel):
    requests: List[KeysExtractRequest] = Field(
        ...,
        min_length=1,
        max_length=50,
        description="List of extraction requests (1-50) to process concurrently",
    )


class KeysExtractBatchResponse(BaseModel):
    success: bool
    time_ms: float
    results: List[KeysExtractResponse]


def _error_response(start: float, message: str) -> KeysExtractResponse:
    return KeysExtractResponse(
        success=False,
        time_ms=round((time.perf_counter() - start) * 1000, 3),
        data={},
        error_message=message,
    )


async def _extract_single(body: KeysExtractRequest) -> KeysExtractResponse:
    start = time.perf_counter()
    has_key_names = bool(body.key_names)
    has_query = body.query is not None

    if has_key_names and has_query:
        return _error_response(start, "provide either `key_names` or `query`, not both")
    if not has_key_names and not has_query:
        return _error_response(start, "either `key_names` or `query` must be provided")
    if not isinstance(body.data, (dict, list)):
        return _error_response(start, "`data` must be a JSON object or array")

    try:
        loop = asyncio.get_running_loop()
        results = await loop.run_in_executor(
            thread_pool,
            KeysExtractor(body.data, body.key_names, body.result_limit, body.query).extract,
        )
    except (TypeError, ValueError) as exc:
        return _error_response(start, str(exc))
    except Exception as exc:
        return _error_response(start, f"extraction failed: {exc}")

    elapsed = round((time.perf_counter() - start) * 1000, 3)
    return KeysExtractResponse(
        success=True,
        time_ms=elapsed,
        data=results,
        error_message=None,
    )


@router.post(
    "/keys/extract",
    response_model=KeysExtractBatchResponse,
    tags=["Keys Extractor"],
    summary="Extract values from nested JSON objects (batch of up to 50 requests processed concurrently)",
)
async def extract_keys_batch(body: KeysExtractBatchRequest):
    """Process up to 50 extraction requests concurrently.

    Each request supports two modes:

    * **key_names** (default) -- recursively walk any JSON object/array and return
      every value attached to the supplied key names, regardless of how deeply
      nested they are. Returns ``{key_name: [values...]}``.

    * **query** -- run an extended JSONPath expression and return each match with
      its full path inside the source document. Returns
      ``{"matches": [{"path": "users.[0].role", "value": "admin"}, ...]}``.
    """
    start = time.perf_counter()
    logger.info("keys_extract_batch | start count=%d", len(body.requests))
    tasks = [_extract_single(req) for req in body.requests]
    results = await asyncio.gather(*tasks)
    elapsed = round((time.perf_counter() - start) * 1000, 3)
    all_ok = all(r.success for r in results)
    logger.info(
        "keys_extract_batch | done total=%d succeeded=%d time_ms=%s",
        len(results),
        sum(1 for r in results if r.success),
        elapsed,
    )
    return KeysExtractBatchResponse(success=all_ok, time_ms=elapsed, results=results)