Fix types, fix tag_queries

PostHog · Jan 12, 2024 · db88128 · db88128
1 parent 0a185f2
commit db88128
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 9 deletions.
diff --git a/posthog/api/mixins.py b/posthog/api/mixins.py
@@ -1,4 +1,4 @@
-from typing import TypeVar, Type, Generic
+from typing import TypeVar, Type
 
 from pydantic import BaseModel, ValidationError
 
@@ -7,7 +7,7 @@
 T = TypeVar("T", bound=BaseModel)
 
 
-class PydanticModelMixin(Generic[T]):
+class PydanticModelMixin:
     def get_model(self, data: dict, model: Type[T]) -> T:
         try:
             return model.model_validate(data)

diff --git a/posthog/api/query.py b/posthog/api/query.py
@@ -77,6 +77,7 @@ def create(self, request, *args, **kwargs) -> Response:
             )
             return Response(query_status.model_dump())
 
+        tag_queries(query=request.data["query"])
         try:
             result = process_query_model(self.team, data.query, refresh_requested=data.refresh)
             return Response(result)
@@ -87,12 +88,6 @@ def create(self, request, *args, **kwargs) -> Response:
             capture_exception(e)
             raise e
 
-    @extend_schema(
-        description="(Experimental)",
-        responses={
-            200: OpenApiResponse(description="Query status"),
-        },
-    )
     @extend_schema(
         description="(Experimental)",
         responses={
@@ -105,6 +100,9 @@ def retrieve(self, request: Request, pk=None, *args, **kwargs) -> JsonResponse:
 
     @extend_schema(
         description="(Experimental)",
+        responses={
+            204: OpenApiResponse(description="Query cancelled"),
+        },
     )
     def destroy(self, request, pk=None, *args, **kwargs):
         cancel_query(self.team.pk, pk)

diff --git a/posthog/api/services/query.py b/posthog/api/services/query.py
@@ -53,6 +53,7 @@ def process_query(
     refresh_requested: Optional[bool] = False,
 ) -> dict:
     model = QuerySchemaRoot.model_validate(query_json)
+    tag_queries(query=query_json)
     return process_query_model(
         team,
         model.root,
@@ -67,7 +68,6 @@ def process_query_model(
     limit_context: Optional[LimitContext] = None,
     refresh_requested: Optional[bool] = False,
 ) -> dict:
-    tag_queries(query=query.kind)  # TODO: ?
     result: dict | BaseModel
 
     if isinstance(query, QUERY_WITH_RUNNER):  # type: ignore