load hog metadata

frontend/src/lib/monaco/CodeEditor.tsx

@@ -1,7 +1,7 @@
 import './CodeEditor.scss'
 
 import MonacoEditor, { type EditorProps, Monaco } from '@monaco-editor/react'
-import { useValues } from 'kea'
+import { useMountedLogic, useValues } from 'kea'
 import { Spinner } from 'lib/lemon-ui/Spinner'
 import { codeEditorLogic } from 'lib/monaco/codeEditorLogic'
 import { hogQLAutocompleteProvider } from 'lib/monaco/hogQLAutocompleteProvider'
@@ -15,25 +15,27 @@ import { useEffect, useRef, useState } from 'react'
 import { themeLogic } from '~/layout/navigation-3000/themeLogic'
 
 export interface CodeEditorProps extends Omit<EditorProps, 'loading' | 'theme'> {
-    logicKey?: string
+    queryKey?: string
 }
 let codeEditorIndex = 0
 
-export function CodeEditor({ logicKey, options, onMount, value, ...editorProps }: CodeEditorProps): JSX.Element {
+export function CodeEditor({ queryKey, options, onMount, value, ...editorProps }: CodeEditorProps): JSX.Element {
     const { isDarkModeOn } = useValues(themeLogic)
     const scrollbarRendering = !inStorybookTestRunner() ? 'auto' : 'hidden'
-    const [realKey] = useState(() => codeEditorIndex++)
     const [monacoAndEditor, setMonacoAndEditor] = useState(
         null as [Monaco, importedEditor.IStandaloneCodeEditor] | null
     )
     const [monaco, editor] = monacoAndEditor ?? []
+
+    const [realKey] = useState(() => codeEditorIndex++)
     const builtCodeEditorLogic = codeEditorLogic({
-        key: logicKey ?? `new/${realKey}`,
+        key: queryKey ?? `new/${realKey}`,
         query: value ?? '',
         language: editorProps.language,
         monaco: monaco,
         editor: editor,
     })
+    useMountedLogic(builtCodeEditorLogic)
 
     // Using useRef, not useState, as we don't want to reload the component when this changes.
     const monacoDisposables = useRef([] as IDisposable[])
@@ -75,6 +77,9 @@ export function CodeEditor({ logicKey, options, onMount, value, ...editorProps }
                         monaco.languages.setLanguageConfiguration('hog', hog.conf)
                         monaco.languages.setMonarchTokensProvider('hog', hog.language)
                     }
+                    monacoDisposables.current.push(
+                        monaco.languages.registerCodeActionProvider('hog', hogQLMetadataProvider(builtCodeEditorLogic))
+                    )
                 }
                 if (editorProps?.language === 'hogql') {
                     if (!monaco.languages.getLanguages().some(({ id }) => id === 'hogql')) {

frontend/src/lib/monaco/codeEditorLogic.tsx

@@ -1,5 +1,5 @@
 import type { Monaco } from '@monaco-editor/react'
-import { actions, kea, key, path, props, propsChanged, reducers, selectors } from 'kea'
+import { actions, kea, key, path, props, propsChanged, selectors } from 'kea'
 import { loaders } from 'kea-loaders'
 // Note: we can oly import types and not values from monaco-editor, because otherwise some Monaco code breaks
 // auto reload in development. Specifically, on this line:
@@ -34,36 +34,25 @@ export const codeEditorLogic = kea<codeEditorLogicType>([
     path(['lib', 'monaco', 'hogQLMetadataProvider']),
     props({} as CodeEditorLogicProps),
     key((props) => props.key),
-    propsChanged(({ actions, props }, oldProps) => {
-        if (props.query !== oldProps.query || props.editor !== oldProps.editor) {
-            actions.setQuery(props.query)
-        }
-    }),
     actions({
-        setQuery: (query: string) => ({ query }),
+        reloadMetadata: true,
     }),
-    reducers(({ props }) => ({
-        query: [props.query, { setQuery: (_, { query }) => query }],
-    })),
-    loaders(({ props, values }) => ({
+    loaders(({ props }) => ({
         metadata: [
             null as null | [string, HogQLMetadataResponse],
             {
-                setQuery: async (_, breakpoint) => {
-                    if (!props.editor || !props.monaco) {
-                        return null
-                    }
+                reloadMetadata: async (_, breakpoint) => {
                     const model = props.editor?.getModel()
-                    if (!model) {
+                    if (!model || !props.monaco || (props.language !== 'hogql' && props.language !== 'hog')) {
                         return null
                     }
                     await breakpoint(300)
-                    const { query } = values
-                    const response = await performQuery<HogQLMetadata>({
-                        kind: NodeKind.HogQLMetadata,
-                        select: query,
-                        filters: props.metadataFilters,
-                    })
+                    const query = props.query
+                    const response = await performQuery<HogQLMetadata>(
+                        props.language === 'hogql'
+                            ? { kind: NodeKind.HogQLMetadata, select: query, filters: props.metadataFilters }
+                            : { kind: NodeKind.HogQLMetadata, program: query, filters: props.metadataFilters }
+                    )
                     breakpoint()
                     return [query, response]
                 },
@@ -72,7 +61,7 @@ export const codeEditorLogic = kea<codeEditorLogicType>([
         modelMarkers: [
             [] as ModelMarker[],
             {
-                setQuerySuccess: ({ metadata }) => {
+                reloadMetadataSuccess: ({ metadata }) => {
                     const model = props.editor?.getModel()
                     if (!model || !metadata) {
                         return []
@@ -128,4 +117,13 @@ export const codeEditorLogic = kea<codeEditorLogicType>([
             },
         ],
     }),
+    propsChanged(({ actions, props }, oldProps) => {
+        if (
+            props.query !== oldProps.query ||
+            props.language !== oldProps.language ||
+            props.editor !== oldProps.editor
+        ) {
+            actions.reloadMetadata()
+        }
+    }),
 ])

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

@@ -157,7 +157,7 @@ export function HogQLQueryEditor(props: HogQLQueryEditorProps): JSX.Element {
                     {/* eslint-disable-next-line react/forbid-dom-props */}
                     <div ref={editorRef} className="resize-y overflow-hidden" style={{ height: EDITOR_HEIGHT }}>
                         <CodeEditor
-                            logicKey={codeEditorKey}
+                            queryKey={codeEditorKey}
                             className="border rounded overflow-hidden h-full"
                             language="hogql"
                             value={queryInput}

frontend/src/queries/schema.json

@@ -4167,6 +4167,10 @@
                     "$ref": "#/definitions/HogQLQueryModifiers",
                     "description": "Modifiers used when performing the query"
                 },
+                "program": {
+                    "description": "Full Hog program",
+                    "type": "string"
+                },
                 "response": {
                     "$ref": "#/definitions/HogQLMetadataResponse"
                 },

frontend/src/queries/schema.ts

@@ -345,6 +345,8 @@ export interface HogQLAutocompleteResponse {
 
 export interface HogQLMetadata extends DataNode<HogQLMetadataResponse> {
     kind: NodeKind.HogQLMetadata
+    /** Full Hog program */
+    program?: string
     /** Full select query to validate (use `select` or `expr`, but not both) */
     select?: string
     /** HogQL expression to validate (use `select` or `expr`, but not both) */

frontend/src/scenes/debug/HogDebug.tsx

@@ -15,8 +15,11 @@ import { HogQuery, HogQueryResponse } from '~/queries/schema'
 export interface HogQueryEditorProps {
     query: HogQuery
     setQuery?: (query: HogQuery) => void
+    queryKey?: string
 }
 
+let uniqueNode = 0
+
 export function HogQueryEditor(props: HogQueryEditorProps): JSX.Element {
     // Using useRef, not useState, as we don't want to reload the component when this changes.
     const monacoDisposables = useRef([] as IDisposable[])
@@ -29,6 +32,7 @@ export function HogQueryEditor(props: HogQueryEditorProps): JSX.Element {
     useEffect(() => {
         setQueryInput(props.query?.code)
     }, [props.query?.code])
+    const [realKey] = useState(() => uniqueNode++)
 
     function saveQuery(): void {
         if (props.setQuery) {
@@ -43,6 +47,7 @@ export function HogQueryEditor(props: HogQueryEditorProps): JSX.Element {
                     {/* eslint-disable-next-line react/forbid-dom-props */}
                     <div className="resize-y overflow-hidden" style={{ height: 222 }}>
                         <CodeEditor
+                            queryKey={props.queryKey ?? `new/${realKey}`}
                             className="border rounded overflow-hidden h-full"
                             language="hog"
                             value={queryInput}
@@ -111,7 +116,7 @@ export function HogDebug({ query, setQuery, queryKey, debug }: HogDebugProps): J
             <div className="space-y-2">
                 {setQuery ? (
                     <>
-                        <HogQueryEditor query={query} setQuery={setQuery} />
+                        <HogQueryEditor query={query} setQuery={setQuery} queryKey={queryKey} />
                         <LemonDivider className="my-4" />
                         <div className="flex gap-2">
                             <Reload />

posthog/hogql/metadata.py

@@ -1,9 +1,11 @@
 from django.conf import settings
+
+from posthog.hogql.bytecode import create_bytecode
 from posthog.hogql.context import HogQLContext
 from posthog.hogql.errors import ExposedHogQLError
 from posthog.hogql.filters import replace_filters
 from posthog.hogql.hogql import translate_hogql
-from posthog.hogql.parser import parse_select
+from posthog.hogql.parser import parse_select, parse_program
 from posthog.hogql.printer import print_ast
 from posthog.hogql.query import create_default_modifiers_for_team
 from posthog.hogql_queries.query_runner import get_query_runner
@@ -54,6 +56,9 @@ def get_hogql_metadata(
                 context=context,
                 dialect="clickhouse",
             )
+        elif isinstance(query.program, str):
+            program = parse_program(query.program)
+            create_bytecode(program, supported_functions={"fetch"}, args=[])
         else:
             raise ValueError("Either expr or select must be provided")
         response.warnings = context.warnings

posthog/schema.py

@@ -4446,6 +4446,7 @@ class HogQLMetadata(BaseModel):
     modifiers: Optional[HogQLQueryModifiers] = Field(
         default=None, description="Modifiers used when performing the query"
     )
+    program: Optional[str] = Field(default=None, description="Full Hog program")
     response: Optional[HogQLMetadataResponse] = None
     select: Optional[str] = Field(
         default=None, description="Full select query to validate (use `select` or `expr`, but not both)"