diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index 3339c51d43e3d..5a9f9216aef9f 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -2225,7 +2225,7 @@ const api = { async get(viewId: DataWarehouseSavedQuery['id']): Promise { return await new ApiRequest().dataWarehouseSavedQuery(viewId).get() }, - async create(data: Partial): Promise { + async create(data: Partial & { types: string[][] }): Promise { return await new ApiRequest().dataWarehouseSavedQueries().create({ data }) }, async delete(viewId: DataWarehouseSavedQuery['id']): Promise { @@ -2233,7 +2233,7 @@ const api = { }, async update( viewId: DataWarehouseSavedQuery['id'], - data: Partial + data: Partial & { types: string[][] } ): Promise { return await new ApiRequest().dataWarehouseSavedQuery(viewId).update({ data }) }, diff --git a/frontend/src/queries/nodes/DataVisualization/DataVisualization.tsx b/frontend/src/queries/nodes/DataVisualization/DataVisualization.tsx index 9e3fbf307c2d4..6a893d2b386f2 100644 --- a/frontend/src/queries/nodes/DataVisualization/DataVisualization.tsx +++ b/frontend/src/queries/nodes/DataVisualization/DataVisualization.tsx @@ -197,6 +197,7 @@ function InternalDataTableVisualization(props: DataTableVisualizationProps): JSX <> void embedded?: boolean editorFooter?: (hasErrors: boolean, errors: string | null, isValidView: boolean) => JSX.Element + queryResponse?: Record } let uniqueNode = 0 @@ -60,6 +61,7 @@ export function HogQLQueryEditor(props: HogQLQueryEditorProps): JSX.Element { key, editor, monaco, + queryResponse: props.queryResponse, } const logic = hogQLQueryEditorLogic(hogQLQueryEditorLogicProps) const { queryInput, prompt, aiAvailable, promptError, promptLoading, multitab } = useValues(logic) diff --git a/frontend/src/queries/nodes/HogQLQuery/hogQLQueryEditorLogic.tsx b/frontend/src/queries/nodes/HogQLQuery/hogQLQueryEditorLogic.tsx index ac798bd39056e..0efaa396e9c94 100644 --- a/frontend/src/queries/nodes/HogQLQuery/hogQLQueryEditorLogic.tsx +++ b/frontend/src/queries/nodes/HogQLQuery/hogQLQueryEditorLogic.tsx @@ -31,6 +31,7 @@ export interface HogQLQueryEditorLogicProps { monaco?: Monaco | null editor?: editor.IStandaloneCodeEditor | null metadataSource?: DataNode + queryResponse?: Record } export const hogQLQueryEditorLogic = kea([ @@ -139,10 +140,13 @@ export const hogQLQueryEditorLogic = kea([ kind: NodeKind.HogQLQuery, query: values.queryInput, } - await dataWarehouseViewsLogic.asyncActions.createDataWarehouseSavedQuery({ name, query }) + const types = props.queryResponse?.types ?? [] + + await dataWarehouseViewsLogic.asyncActions.createDataWarehouseSavedQuery({ name, query, types }) }, onUpdateView: async () => { - actions.updateView(values.queryInput) + const types = props.queryResponse?.types ?? [] + actions.updateView(values.queryInput, types) }, })), ]) diff --git a/frontend/src/scenes/data-warehouse/editor/ResultPane.tsx b/frontend/src/scenes/data-warehouse/editor/ResultPane.tsx index ac7915651296a..6f12365999db5 100644 --- a/frontend/src/scenes/data-warehouse/editor/ResultPane.tsx +++ b/frontend/src/scenes/data-warehouse/editor/ResultPane.tsx @@ -104,6 +104,7 @@ export function ResultPane({ kind: NodeKind.HogQLQuery, query: queryInput, }, + types: response?.types ?? [], }) } > diff --git a/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx b/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx index cad1c656c0b91..ef85ff1e740bb 100644 --- a/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx +++ b/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx @@ -322,7 +322,18 @@ export const multitabEditorLogic = kea([ kind: NodeKind.HogQLQuery, query: values.queryInput, } - await dataWarehouseViewsLogic.asyncActions.createDataWarehouseSavedQuery({ name, query }) + + const logic = dataNodeLogic({ + key: values.activeTabKey, + query: { + kind: NodeKind.HogQLQuery, + query: values.queryInput, + }, + }) + + const types = logic.values.response?.types ?? [] + + await dataWarehouseViewsLogic.asyncActions.createDataWarehouseSavedQuery({ name, query, types }) }, reloadMetadata: async (_, breakpoint) => { const model = props.editor?.getModel() diff --git a/frontend/src/scenes/data-warehouse/saved_queries/dataWarehouseViewsLogic.tsx b/frontend/src/scenes/data-warehouse/saved_queries/dataWarehouseViewsLogic.tsx index 37c744e633d9b..d66a0285526ba 100644 --- a/frontend/src/scenes/data-warehouse/saved_queries/dataWarehouseViewsLogic.tsx +++ b/frontend/src/scenes/data-warehouse/saved_queries/dataWarehouseViewsLogic.tsx @@ -35,7 +35,9 @@ export const dataWarehouseViewsLogic = kea([ } return savedQueries.results }, - createDataWarehouseSavedQuery: async (view: Partial) => { + createDataWarehouseSavedQuery: async ( + view: Partial & { types: string[][] } + ) => { const newView = await api.dataWarehouseSavedQueries.create(view) lemonToast.success(`${newView.name ?? 'View'} successfully created`) @@ -46,7 +48,9 @@ export const dataWarehouseViewsLogic = kea([ await api.dataWarehouseSavedQueries.delete(viewId) return values.dataWarehouseSavedQueries.filter((view) => view.id !== viewId) }, - updateDataWarehouseSavedQuery: async (view: Partial & { id: string }) => { + updateDataWarehouseSavedQuery: async ( + view: Partial & { id: string; types: string[][] } + ) => { const newView = await api.dataWarehouseSavedQueries.update(view.id, view) return values.dataWarehouseSavedQueries.map((savedQuery) => { if (savedQuery.id === view.id) { diff --git a/frontend/src/scenes/data-warehouse/settings/dataWarehouseSceneLogic.ts b/frontend/src/scenes/data-warehouse/settings/dataWarehouseSceneLogic.ts index 85e8283b57fad..c3f05d5562b76 100644 --- a/frontend/src/scenes/data-warehouse/settings/dataWarehouseSceneLogic.ts +++ b/frontend/src/scenes/data-warehouse/settings/dataWarehouseSceneLogic.ts @@ -11,6 +11,7 @@ import { urls } from 'scenes/urls' import { DatabaseSchemaMaterializedViewTable, DatabaseSchemaTable, + DatabaseSchemaViewTable, DatabaseSerializedFieldType, HogQLQuery, NodeKind, @@ -55,7 +56,7 @@ export const dataWarehouseSceneLogic = kea([ deleteDataWarehouseTable: (tableId: string) => ({ tableId }), toggleSchemaModal: true, setEditingView: (id: string | null) => ({ id }), - updateView: (query: string) => ({ query }), + updateView: (query: string, types: string[][]) => ({ query, types }), })), reducers({ selectedRow: [ @@ -278,16 +279,17 @@ export const dataWarehouseSceneLogic = kea([ }) } }, - updateView: ({ query }) => { + updateView: ({ query, types }) => { if (values.editingView) { const newViewQuery: HogQLQuery = { kind: NodeKind.HogQLQuery, query: query, } const oldView = values.viewsMapById[values.editingView] - const newView = { + const newView: DatabaseSchemaViewTable & { types: string[][] } = { ...oldView, query: newViewQuery, + types, } actions.updateDataWarehouseSavedQuery(newView) } diff --git a/posthog/hogql/constants.py b/posthog/hogql/constants.py index 5e64111632997..9a7c463ba3ab0 100644 --- a/posthog/hogql/constants.py +++ b/posthog/hogql/constants.py @@ -1,5 +1,6 @@ from datetime import date, datetime from enum import StrEnum +import sys from typing import Optional, Literal, TypeAlias from uuid import UUID from pydantic import ConfigDict, BaseModel @@ -53,6 +54,7 @@ class LimitContext(StrEnum): EXPORT = "export" COHORT_CALCULATION = "cohort_calculation" HEATMAPS = "heatmaps" + SAVED_QUERY = "saved_query" def get_max_limit_for_context(limit_context: LimitContext) -> int: @@ -62,6 +64,8 @@ def get_max_limit_for_context(limit_context: LimitContext) -> int: return MAX_SELECT_HEATMAPS_LIMIT # 1M elif limit_context == LimitContext.COHORT_CALCULATION: return MAX_SELECT_COHORT_CALCULATION_LIMIT # 1b + elif limit_context == LimitContext.SAVED_QUERY: + return sys.maxsize # Max python int else: raise ValueError(f"Unexpected LimitContext value: {limit_context}") @@ -76,6 +80,8 @@ def get_default_limit_for_context(limit_context: LimitContext) -> int: return MAX_SELECT_HEATMAPS_LIMIT # 1M elif limit_context == LimitContext.COHORT_CALCULATION: return MAX_SELECT_COHORT_CALCULATION_LIMIT # 1b + elif limit_context == LimitContext.SAVED_QUERY: + return sys.maxsize # Max python int else: raise ValueError(f"Unexpected LimitContext value: {limit_context}") diff --git a/posthog/temporal/data_modeling/run_workflow.py b/posthog/temporal/data_modeling/run_workflow.py index f6f35bb9a67b3..eab0782b6c4fd 100644 --- a/posthog/temporal/data_modeling/run_workflow.py +++ b/posthog/temporal/data_modeling/run_workflow.py @@ -23,7 +23,7 @@ from django.conf import settings from dlt.common.libs.deltalake import get_delta_tables -from posthog.hogql.constants import HogQLGlobalSettings +from posthog.hogql.constants import HogQLGlobalSettings, LimitContext from posthog.hogql.database.database import create_hogql_database from posthog.hogql.query import execute_hogql_query from posthog.models import Team @@ -347,7 +347,9 @@ def hogql_table(query: str, team: Team, table_name: str, table_columns: dlt_typi async def get_hogql_rows(): settings = HogQLGlobalSettings(max_execution_time=60 * 10) # 10 mins, same as the /query endpoint async workers - response = await asyncio.to_thread(execute_hogql_query, query, team, settings=settings) + response = await asyncio.to_thread( + execute_hogql_query, query, team, settings=settings, limit_context=LimitContext.SAVED_QUERY + ) if not response.columns: raise EmptyHogQLResponseColumnsError() diff --git a/posthog/warehouse/api/saved_query.py b/posthog/warehouse/api/saved_query.py index 2d8ef156aa6b0..764998c264074 100644 --- a/posthog/warehouse/api/saved_query.py +++ b/posthog/warehouse/api/saved_query.py @@ -19,7 +19,13 @@ from posthog.hogql.printer import print_ast from posthog.temporal.common.client import sync_connect from posthog.temporal.data_modeling.run_workflow import RunWorkflowInputs, Selector -from posthog.warehouse.models import DataWarehouseJoin, DataWarehouseModelPath, DataWarehouseSavedQuery +from posthog.warehouse.models import ( + CLICKHOUSE_HOGQL_MAPPING, + DataWarehouseJoin, + DataWarehouseModelPath, + DataWarehouseSavedQuery, + clean_type, +) import uuid @@ -70,7 +76,20 @@ def create(self, validated_data): view = DataWarehouseSavedQuery(**validated_data) # The columns will be inferred from the query try: - view.columns = view.get_columns() + client_types = self.context["request"].data.get("types", []) + if len(client_types) == 0: + view.columns = view.get_columns() + else: + columns = { + str(item[0]): { + "hogql": CLICKHOUSE_HOGQL_MAPPING[clean_type(str(item[1]))].__name__, + "clickhouse": item[1], + "valid": True, + } + for item in client_types + } + view.columns = columns + view.external_tables = view.s3_tables except Exception as err: raise serializers.ValidationError(str(err)) diff --git a/posthog/warehouse/api/test/test_saved_query.py b/posthog/warehouse/api/test/test_saved_query.py index 6bc7f0c07ac23..6790887374873 100644 --- a/posthog/warehouse/api/test/test_saved_query.py +++ b/posthog/warehouse/api/test/test_saved_query.py @@ -1,7 +1,8 @@ +from unittest.mock import patch import uuid from posthog.test.base import APIBaseTest -from posthog.warehouse.models import DataWarehouseModelPath +from posthog.warehouse.models import DataWarehouseModelPath, DataWarehouseSavedQuery class TestSavedQuery(APIBaseTest): @@ -34,6 +35,36 @@ def test_create(self): ], ) + def test_create_with_types(self): + with patch.object(DataWarehouseSavedQuery, "get_columns") as mock_get_columns: + response = self.client.post( + f"/api/projects/{self.team.id}/warehouse_saved_queries/", + { + "name": "event_view", + "query": { + "kind": "HogQLQuery", + "query": "select event as event from events LIMIT 100", + }, + "types": [["event", "Nullable(String)"]], + }, + ) + assert response.status_code == 201 + saved_query = response.json() + assert saved_query["name"] == "event_view" + assert saved_query["columns"] == [ + { + "key": "event", + "name": "event", + "type": "string", + "schema_valid": True, + "fields": None, + "table": None, + "chain": None, + } + ] + + mock_get_columns.assert_not_called() + def test_create_name_overlap_error(self): response = self.client.post( f"/api/projects/{self.team.id}/warehouse_saved_queries/",