chore(data-warehouse): Allowed saving of views without querying again

frontend/src/lib/api.ts

@@ -2230,15 +2230,15 @@ const api = {
         async get(viewId: DataWarehouseSavedQuery['id']): Promise<DataWarehouseSavedQuery> {
             return await new ApiRequest().dataWarehouseSavedQuery(viewId).get()
         },
-        async create(data: Partial<DataWarehouseSavedQuery>): Promise<DataWarehouseSavedQuery> {
+        async create(data: Partial<DataWarehouseSavedQuery> & { types: string[][] }): Promise<DataWarehouseSavedQuery> {
             return await new ApiRequest().dataWarehouseSavedQueries().create({ data })
         },
         async delete(viewId: DataWarehouseSavedQuery['id']): Promise<void> {
             await new ApiRequest().dataWarehouseSavedQuery(viewId).delete()
         },
         async update(
             viewId: DataWarehouseSavedQuery['id'],
-            data: Partial<DataWarehouseSavedQuery>
+            data: Partial<DataWarehouseSavedQuery> & { types: string[][] }
         ): Promise<DataWarehouseSavedQuery> {
             return await new ApiRequest().dataWarehouseSavedQuery(viewId).update({ data })
         },

frontend/src/queries/nodes/DataVisualization/DataVisualization.tsx

@@ -188,6 +188,7 @@ function InternalDataTableVisualization(props: DataTableVisualizationProps): JSX
                     <>
                         <HogQLQueryEditor
                             query={query.source}
+                            queryResponse={response ?? undefined}
                             setQuery={setQuerySource}
                             embedded
                             onChange={setEditorQuery}

frontend/src/queries/nodes/HogQLQuery/HogQLQueryEditor.tsx

@@ -30,6 +30,7 @@ export interface HogQLQueryEditorProps {
     onChange?: (query: string) => void
     embedded?: boolean
     editorFooter?: (hasErrors: boolean, errors: string | null, isValidView: boolean) => JSX.Element
+    queryResponse?: Record<string, any>
 }
 
 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)

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<string, any>
 }
 
 export const hogQLQueryEditorLogic = kea<hogQLQueryEditorLogicType>([
@@ -139,10 +140,13 @@ export const hogQLQueryEditorLogic = kea<hogQLQueryEditorLogicType>([
                 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)
         },
     })),
 ])

frontend/src/scenes/data-warehouse/saved_queries/dataWarehouseViewsLogic.tsx

@@ -35,7 +35,9 @@ export const dataWarehouseViewsLogic = kea<dataWarehouseViewsLogicType>([
                     }
                     return savedQueries.results
                 },
-                createDataWarehouseSavedQuery: async (view: Partial<DatabaseSchemaViewTable>) => {
+                createDataWarehouseSavedQuery: async (
+                    view: Partial<DatabaseSchemaViewTable> & { types: string[][] }
+                ) => {
                     const newView = await api.dataWarehouseSavedQueries.create(view)
 
                     lemonToast.success(`${newView.name ?? 'View'} successfully created`)
@@ -46,7 +48,9 @@ export const dataWarehouseViewsLogic = kea<dataWarehouseViewsLogicType>([
                     await api.dataWarehouseSavedQueries.delete(viewId)
                     return values.dataWarehouseSavedQueries.filter((view) => view.id !== viewId)
                 },
-                updateDataWarehouseSavedQuery: async (view: Partial<DatabaseSchemaViewTable> & { id: string }) => {
+                updateDataWarehouseSavedQuery: async (
+                    view: Partial<DatabaseSchemaViewTable> & { id: string; types: string[][] }
+                ) => {
                     const newView = await api.dataWarehouseSavedQueries.update(view.id, view)
                     return values.dataWarehouseSavedQueries.map((savedQuery) => {
                         if (savedQuery.id === view.id) {

frontend/src/scenes/data-warehouse/settings/dataWarehouseSceneLogic.ts

@@ -11,6 +11,7 @@
 import {
     DatabaseSchemaMaterializedViewTable,
     DatabaseSchemaTable,
+    DatabaseSchemaViewTable,
     DatabaseSerializedFieldType,
     HogQLQuery,
     NodeKind,
@@ -55,7 +56,7 @@
         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,18 +279,19 @@
                 })
             }
         },
-        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 = {
                     ...oldView,
                     query: newViewQuery,
+                    types,
                 }
                 actions.updateDataWarehouseSavedQuery(newView)
             }
         },
     })),

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}")
 

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()

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))

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/",