feat(hogql): Allow lazy joins on lazy tables with requested fields

cypress/e2e/auth.cy.ts

@@ -84,7 +84,7 @@ describe('Auth', () => {
         cy.visit('/signup')
         cy.location('pathname').should('eq', '/project/1')
     })
-    
+
     it('Logout in another tab results in logout in the current tab too', () => {
         cy.window().then(async (win) => {
             // Hit /logout *in the background* by using fetch()

cypress/productAnalytics/index.ts

@@ -209,14 +209,14 @@ export const dashboard = {
         cy.get('[data-attr="prop-val-0"]').click({ force: true })
         cy.get('.PropertyFilterButton').should('have.length', 1)
     },
-    addPropertyFilter(type: string = "Browser", value: string = "Chrome"): void {
+    addPropertyFilter(type: string = 'Browser', value: string = 'Chrome'): void {
         cy.get('.PropertyFilterButton').should('have.length', 0)
         cy.get('[data-attr="property-filter-0"]').click()
-        cy.get('[data-attr="taxonomic-filter-searchfield"]').click().type("Browser").wait(1000)
+        cy.get('[data-attr="taxonomic-filter-searchfield"]').click().type('Browser').wait(1000)
         cy.get('[data-attr="prop-filter-event_properties-0"]').click({ force: true })
         cy.get('.ant-select-selector').type(value)
         cy.get('.ant-select-item-option-content').click({ force: true })
-    }
+    },
 }
 
 export function createInsight(insightName: string): void {

frontend/src/lib/api.ts

@@ -1926,7 +1926,13 @@ const api = {
             viewId: DataWarehouseViewLink['id'],
             data: Pick<
                 DataWarehouseViewLink,
-                'source_table_name' | 'source_table_key' | 'joining_table_name' | 'joining_table_key' | 'field_name'
+                | 'source_table_name'
+                | 'source_table_key'
+                | 'joining_table_name'
+                | 'joining_table_key'
+                | 'field_name'
+                | 'source_table_key_hogql'
+                | 'joining_table_key_hogql'
             >
         ): Promise<DataWarehouseViewLink> {
             return await new ApiRequest().dataWarehouseViewLink(viewId).update({ data })

frontend/src/scenes/data-warehouse/ViewLinkModal.tsx

@@ -1,12 +1,22 @@
 import './ViewLinkModal.scss'
 
 import { IconTrash } from '@posthog/icons'
-import { LemonButton, LemonDivider, LemonInput, LemonModal, LemonSelect, LemonTag } from '@posthog/lemon-ui'
+import {
+    LemonButton,
+    LemonDivider,
+    LemonDropdown,
+    LemonInput,
+    LemonModal,
+    LemonSelect,
+    LemonTag,
+} from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { Field, Form } from 'kea-forms'
 import { CodeSnippet, Language } from 'lib/components/CodeSnippet'
+import { HogQLEditor } from 'lib/components/HogQLEditor/HogQLEditor'
 import { IconSwapHoriz } from 'lib/lemon-ui/icons'
-import { viewLinkLogic } from 'scenes/data-warehouse/viewLinkLogic'
+import { useState } from 'react'
+import { HOGQL_IDENTIFIER, viewLinkLogic } from 'scenes/data-warehouse/viewLinkLogic'
 
 import { DatabaseSchemaQueryResponseField } from '~/queries/schema'
 
@@ -44,8 +54,21 @@ export function ViewLinkForm(): JSX.Element {
         error,
         fieldName,
         isNewJoin,
+        selectedSourceKey,
+        selectedSourceKeyHogQL,
+        selectedJoiningKey,
+        selectedJoiningKeyHogQL,
     } = useValues(viewLinkLogic)
-    const { selectJoiningTable, toggleJoinTableModal, selectSourceTable, setFieldName } = useActions(viewLinkLogic)
+    const {
+        selectJoiningTable,
+        toggleJoinTableModal,
+        selectSourceTable,
+        setFieldName,
+        selectSourceKey,
+        selectSourceKeyHogQL,
+        selectJoiningKey,
+        selectJoiningKeyHogQL,
+    } = useActions(viewLinkLogic)
 
     return (
         <Form logic={viewLinkLogic} formKey="viewLink" enableFormOnSubmit>
@@ -82,12 +105,25 @@ export function ViewLinkForm(): JSX.Element {
                     <div className="w-50">
                         <span className="l4">Source Table Key</span>
                         <Field name="source_table_key">
-                            <LemonSelect
-                                fullWidth
-                                disabledReason={selectedSourceTableName ? '' : 'Select a table to choose join key'}
-                                options={sourceTableKeys}
-                                placeholder="Select a key"
-                            />
+                            <>
+                                <LemonSelect
+                                    fullWidth
+                                    onSelect={selectSourceKey}
+                                    value={selectedSourceKey ?? undefined}
+                                    disabledReason={selectedSourceTableName ? '' : 'Select a table to choose join key'}
+                                    options={[
+                                        ...sourceTableKeys,
+                                        { value: HOGQL_IDENTIFIER, label: <span>HogQL Expression</span> },
+                                    ]}
+                                    placeholder="Select a key"
+                                />
+                                {selectedSourceKey === HOGQL_IDENTIFIER && (
+                                    <HogQLDropdown
+                                        hogQLValue={selectedSourceKeyHogQL ?? ''}
+                                        onHogQLValueChange={selectSourceKeyHogQL}
+                                    />
+                                )}
+                            </>
                         </Field>
                     </div>
                     <div className="mt-5">
@@ -96,12 +132,25 @@ export function ViewLinkForm(): JSX.Element {
                     <div className="w-50">
                         <span className="l4">Joining Table Key</span>
                         <Field name="joining_table_key">
-                            <LemonSelect
-                                fullWidth
-                                disabledReason={selectedJoiningTable ? '' : 'Select a table to choose join key'}
-                                options={joiningTableKeys}
-                                placeholder="Select a key"
-                            />
+                            <>
+                                <LemonSelect
+                                    fullWidth
+                                    onSelect={selectJoiningKey}
+                                    value={selectedJoiningKey ?? undefined}
+                                    disabledReason={selectedJoiningTable ? '' : 'Select a table to choose join key'}
+                                    options={[
+                                        ...joiningTableKeys,
+                                        { value: HOGQL_IDENTIFIER, label: <span>HogQL Expression</span> },
+                                    ]}
+                                    placeholder="Select a key"
+                                />
+                                {selectedJoiningKey === HOGQL_IDENTIFIER && (
+                                    <HogQLDropdown
+                                        hogQLValue={selectedJoiningKeyHogQL ?? ''}
+                                        onHogQLValueChange={selectJoiningKeyHogQL}
+                                    />
+                                )}
+                            </>
                         </Field>
                     </div>
                 </div>
@@ -151,6 +200,47 @@ export function ViewLinkForm(): JSX.Element {
     )
 }
 
+const HogQLDropdown = ({
+    hogQLValue,
+    onHogQLValueChange,
+}: {
+    hogQLValue: string
+    onHogQLValueChange: (hogQLValue: string) => void
+}): JSX.Element => {
+    const [isHogQLDropdownVisible, setIsHogQLDropdownVisible] = useState(false)
+
+    return (
+        <div className="flex-auto overflow-hidden">
+            <LemonDropdown
+                visible={isHogQLDropdownVisible}
+                closeOnClickInside={false}
+                onClickOutside={() => setIsHogQLDropdownVisible(false)}
+                overlay={
+                    // eslint-disable-next-line react/forbid-dom-props
+                    <div className="w-120" style={{ maxWidth: 'max(60vw, 20rem)' }}>
+                        <HogQLEditor
+                            disablePersonProperties
+                            value={hogQLValue}
+                            onChange={(currentValue) => {
+                                onHogQLValueChange(currentValue)
+                                setIsHogQLDropdownVisible(false)
+                            }}
+                        />
+                    </div>
+                }
+            >
+                <LemonButton
+                    fullWidth
+                    type="secondary"
+                    onClick={() => setIsHogQLDropdownVisible(!isHogQLDropdownVisible)}
+                >
+                    <code>{hogQLValue}</code>
+                </LemonButton>
+            </LemonDropdown>
+        </div>
+    )
+}
+
 interface ViewLinkDeleteButtonProps {
     table: string
     column: string

frontend/src/scenes/data-warehouse/viewLinkLogic.tsx

@@ -15,12 +15,12 @@ import { ViewLinkKeyLabel } from './ViewLinkModal'
 const NEW_VIEW_LINK: DataWarehouseViewLink = {
     id: 'new',
     source_table_name: undefined,
-    source_table_key: undefined,
     joining_table_name: undefined,
-    joining_table_key: undefined,
     field_name: undefined,
 }
 
+export const HOGQL_IDENTIFIER = '$hogql'
+
 export interface KeySelectOption {
     value: string
     label: JSX.Element
@@ -37,9 +37,13 @@ export const viewLinkLogic = kea<viewLinkLogicType>([
         ],
         actions: [databaseTableListLogic, ['loadDatabase'], dataWarehouseJoinsLogic, ['loadJoins']],
     }),
-    actions({
+    actions(({ values }) => ({
         selectJoiningTable: (selectedTableName: string) => ({ selectedTableName }),
         selectSourceTable: (selectedTableName: string) => ({ selectedTableName }),
+        selectSourceKey: (selectedKey: string) => ({ selectedKey, sourceTable: values.selectedSourceTable }),
+        selectSourceKeyHogQL: (hogQL: string) => ({ hogQL }),
+        selectJoiningKey: (selectedKey: string) => ({ selectedKey, joiningTable: values.selectedJoiningTable }),
+        selectJoiningKeyHogQL: (hogQL: string) => ({ hogQL }),
         toggleJoinTableModal: true,
         toggleEditJoinModal: (join: DataWarehouseViewLink) => ({ join }),
         toggleNewJoinModal: true,
@@ -48,7 +52,7 @@ export const viewLinkLogic = kea<viewLinkLogicType>([
         setError: (error: string) => ({ error }),
         setFieldName: (fieldName: string) => ({ fieldName }),
         clearModalFields: true,
-    }),
+    })),
     reducers({
         joinToEdit: [
             null as DataWarehouseViewLink | null,
@@ -84,6 +88,54 @@ export const viewLinkLogic = kea<viewLinkLogicType>([
                 clearModalFields: () => null,
             },
         ],
+        selectedSourceKey: [
+            null as string | null,
+            {
+                selectSourceKey: (_, { selectedKey }) => selectedKey,
+                toggleEditJoinModal: (_, { join }) => join.source_table_key ?? null,
+            },
+        ],
+        selectedSourceKeyHogQL: [
+            null as string | null,
+            {
+                selectSourceKeyHogQL: (_, { hogQL }) => hogQL,
+                toggleEditJoinModal: (_, { join }) => join.source_table_key_hogql ?? null,
+                selectSourceKey: (state, { selectedKey, sourceTable }) => {
+                    if (state === null && selectedKey === HOGQL_IDENTIFIER) {
+                        const firstColumn = sourceTable?.columns?.[0]?.key
+                        if (firstColumn) {
+                            return `lower(${firstColumn})`
+                        }
+                    }
+
+                    return state
+                },
+            },
+        ],
+        selectedJoiningKey: [
+            null as string | null,
+            {
+                selectJoiningKey: (_, { selectedKey }) => selectedKey,
+                toggleEditJoinModal: (_, { join }) => join.joining_table_key ?? null,
+            },
+        ],
+        selectedJoiningKeyHogQL: [
+            null as string | null,
+            {
+                selectJoiningKeyHogQL: (_, { hogQL }) => hogQL,
+                toggleEditJoinModal: (_, { join }) => join.joining_table_key_hogql ?? null,
+                selectJoiningKey: (state, { selectedKey, joiningTable }) => {
+                    if (state === null && selectedKey === HOGQL_IDENTIFIER) {
+                        const firstColumn = joiningTable?.columns?.[0]?.key
+                        if (firstColumn) {
+                            return `lower(${firstColumn})`
+                        }
+                    }
+
+                    return state
+                },
+            },
+        ],
         fieldName: [
             '' as string,
             {
@@ -112,44 +164,23 @@ export const viewLinkLogic = kea<viewLinkLogicType>([
     forms(({ actions, values }) => ({
         viewLink: {
             defaults: NEW_VIEW_LINK,
-            errors: ({ source_table_name, joining_table_name, joining_table_key, source_table_key }) => {
-                let joining_table_key_err: string | undefined = undefined
-                let source_table_key_err: string | undefined = undefined
-
-                if (!joining_table_key) {
-                    joining_table_key_err = 'Must select a join key'
-                }
-
-                if (!source_table_key) {
-                    source_table_key_err = 'Must select a join key'
-                }
-
-                if (
-                    joining_table_key &&
-                    source_table_key &&
-                    values.selectedJoiningTable?.columns?.find((n) => n.key == joining_table_key)?.type !==
-                        values.selectedSourceTable?.columns?.find((n) => n.key == source_table_key)?.type
-                ) {
-                    joining_table_key_err = 'Join key types must match'
-                    source_table_key_err = 'Join key types must match'
-                }
-
+            errors: ({ source_table_name, joining_table_name }) => {
                 return {
                     source_table_name: values.isNewJoin && !source_table_name ? 'Must select a table' : undefined,
                     joining_table_name: !joining_table_name ? 'Must select a table' : undefined,
-                    source_table_key: source_table_key_err,
-                    joining_table_key: joining_table_key_err,
                 }
             },
-            submit: async ({ joining_table_name, source_table_name, source_table_key, joining_table_key }) => {
+            submit: async ({ joining_table_name, source_table_name }) => {
                 if (values.joinToEdit?.id && values.selectedSourceTable) {
                     // Edit join
                     try {
                         await api.dataWarehouseViewLinks.update(values.joinToEdit.id, {
                             source_table_name: source_table_name ?? values.selectedSourceTable.name,
-                            source_table_key,
+                            source_table_key: values.selectedSourceKey ?? undefined,
+                            source_table_key_hogql: values.selectedSourceKeyHogQL ?? undefined,
                             joining_table_name,
-                            joining_table_key,
+                            joining_table_key: values.selectedJoiningKey ?? undefined,
+                            joining_table_key_hogql: values.selectedJoiningKeyHogQL ?? undefined,
                             field_name: values.fieldName,
                         })
 
@@ -164,9 +195,11 @@ export const viewLinkLogic = kea<viewLinkLogicType>([
                     try {
                         await api.dataWarehouseViewLinks.create({
                             source_table_name: source_table_name ?? values.selectedSourceTable.name,
-                            source_table_key,
+                            source_table_key: values.selectedSourceKey ?? undefined,
+                            source_table_key_hogql: values.selectedSourceKeyHogQL ?? undefined,
                             joining_table_name,
-                            joining_table_key,
+                            joining_table_key: values.selectedJoiningKey ?? undefined,
+                            joining_table_key_hogql: values.selectedJoiningKeyHogQL ?? undefined,
                             field_name: values.fieldName,
                         })
 

frontend/src/types.ts

@@ -3508,8 +3508,10 @@ export interface DataWarehouseViewLink {
     id: string
     source_table_name?: string
     source_table_key?: string
+    source_table_key_hogql?: string
     joining_table_name?: string
     joining_table_key?: string
+    joining_table_key_hogql?: string
     field_name?: string
     created_by?: UserBasicType | null
     created_at?: string | null

latest_migrations.manifest

@@ -5,7 +5,7 @@ contenttypes: 0002_remove_content_type_name
 ee: 0015_add_verified_properties
 otp_static: 0002_throttling
 otp_totp: 0002_auto_20190420_0723
-posthog: 0395_alter_batchexportbackfill_end_at
+posthog: 0396_datawarehousejoin_joining_table_key_hogql_and_more
 sessions: 0001_initial
 social_django: 0010_uid_db_index
 two_factor: 0007_auto_20201201_1019