From 95798a23ca2afb737a2298933443163c637d6ab6 Mon Sep 17 00:00:00 2001
From: timgl <tim@glsr.nl>
Date: Mon, 20 Jan 2025 14:05:21 -0800
Subject: [PATCH] fix(cohort): Fix creating static cohort from SQL (#27676)

---
 .../src/scenes/insights/InsightPageHeader.tsx | 91 +++++++++----------
 .../src/scenes/insights/insightDataLogic.tsx  | 14 ++-
 2 files changed, 56 insertions(+), 49 deletions(-)

diff --git a/frontend/src/scenes/insights/InsightPageHeader.tsx b/frontend/src/scenes/insights/InsightPageHeader.tsx
index f193c615a1010..bdcd2023c490d 100644
--- a/frontend/src/scenes/insights/InsightPageHeader.tsx
+++ b/frontend/src/scenes/insights/InsightPageHeader.tsx
@@ -242,9 +242,11 @@ export function InsightPageHeader({ insightLogicProps }: { insightLogicProps: In
                                             label="Debug panel"
                                         />
                                     ) : null}
-                                    {hogQL && (
-                                        <>
-                                            <LemonDivider />
+
+                                    {(hogQL || showCohortButton) && <LemonDivider />}
+                                    {hogQL &&
+                                        !isHogQLQuery(query) &&
+                                        !(isDataVisualizationNode(query) && isHogQLQuery(query.source)) && (
                                             <LemonButton
                                                 data-attr="edit-insight-sql"
                                                 onClick={() => {
@@ -263,50 +265,47 @@ export function InsightPageHeader({ insightLogicProps }: { insightLogicProps: In
                                             >
                                                 Edit SQL directly
                                             </LemonButton>
-                                            {showCohortButton && (
-                                                <LemonButton
-                                                    data-attr="edit-insight-sql"
-                                                    onClick={() => {
-                                                        LemonDialog.openForm({
-                                                            title: 'Save as static cohort',
-                                                            description: (
-                                                                <div className="mt-2">
-                                                                    Your query must export a <code>person_id</code>,{' '}
-                                                                    <code>actor_id</code> or <code>id</code> column,
-                                                                    which must match the <code>id</code> of the{' '}
-                                                                    <code>persons</code> table
-                                                                </div>
-                                                            ),
-                                                            initialValues: {
-                                                                name: '',
-                                                            },
-                                                            content: (
-                                                                <LemonField name="name">
-                                                                    <LemonInput
-                                                                        data-attr="insight-name"
-                                                                        placeholder="Name of the new cohort"
-                                                                        autoFocus
-                                                                    />
-                                                                </LemonField>
-                                                            ),
-                                                            errors: {
-                                                                name: (name) =>
-                                                                    !name ? 'You must enter a name' : undefined,
-                                                            },
-                                                            onSubmit: async ({ name }) => {
-                                                                createStaticCohort(name, {
-                                                                    kind: NodeKind.HogQLQuery,
-                                                                    query: hogQL,
-                                                                })
-                                                            },
+                                        )}
+                                    {hogQL && showCohortButton && (
+                                        <LemonButton
+                                            data-attr="edit-insight-sql"
+                                            onClick={() => {
+                                                LemonDialog.openForm({
+                                                    title: 'Save as static cohort',
+                                                    description: (
+                                                        <div className="mt-2">
+                                                            Your query must export a <code>person_id</code>,{' '}
+                                                            <code>actor_id</code> or <code>id</code> column, which must
+                                                            match the <code>id</code> of the <code>persons</code> table
+                                                        </div>
+                                                    ),
+                                                    initialValues: {
+                                                        name: '',
+                                                    },
+                                                    content: (
+                                                        <LemonField name="name">
+                                                            <LemonInput
+                                                                data-attr="insight-name"
+                                                                placeholder="Name of the new cohort"
+                                                                autoFocus
+                                                            />
+                                                        </LemonField>
+                                                    ),
+                                                    errors: {
+                                                        name: (name) => (!name ? 'You must enter a name' : undefined),
+                                                    },
+                                                    onSubmit: async ({ name }) => {
+                                                        createStaticCohort(name, {
+                                                            kind: NodeKind.HogQLQuery,
+                                                            query: hogQL,
                                                         })
-                                                    }}
-                                                    fullWidth
-                                                >
-                                                    Save as static cohort
-                                                </LemonButton>
-                                            )}
-                                        </>
+                                                    },
+                                                })
+                                            }}
+                                            fullWidth
+                                        >
+                                            Save as static cohort
+                                        </LemonButton>
                                     )}
                                     {hasDashboardItemId && (
                                         <>
diff --git a/frontend/src/scenes/insights/insightDataLogic.tsx b/frontend/src/scenes/insights/insightDataLogic.tsx
index ced59ae6c2914..25945a236be33 100644
--- a/frontend/src/scenes/insights/insightDataLogic.tsx
+++ b/frontend/src/scenes/insights/insightDataLogic.tsx
@@ -13,7 +13,7 @@ import { insightVizDataNodeKey } from '~/queries/nodes/InsightViz/InsightViz'
 import { getDefaultQuery, queryFromKind } from '~/queries/nodes/InsightViz/utils'
 import { queryExportContext } from '~/queries/query'
 import { DataVisualizationNode, InsightVizNode, Node, NodeKind } from '~/queries/schema/schema-general'
-import { isDataTableNode, isDataVisualizationNode, isHogQuery, isInsightVizNode } from '~/queries/utils'
+import { isDataTableNode, isDataVisualizationNode, isHogQLQuery, isHogQuery, isInsightVizNode } from '~/queries/utils'
 import { ExportContext, InsightLogicProps, InsightType } from '~/types'
 
 import { teamLogic } from '../teamLogic'
@@ -171,8 +171,16 @@ export const insightDataLogic = kea<insightDataLogicType>([
         ],
 
         hogQL: [
-            (s) => [s.insightData],
-            (insightData): string | null => {
+            (s) => [s.insightData, s.query],
+            (insightData, query): string | null => {
+                // Try to get it from the query itself, so we don't have to wait for the response
+                if (isDataVisualizationNode(query) && isHogQLQuery(query.source)) {
+                    return query.source.query
+                }
+                if (isHogQLQuery(query)) {
+                    return query.query
+                }
+                // Otherwise, get it from the response
                 if (insightData && 'hogql' in insightData && insightData.hogql !== '') {
                     return insightData.hogql
                 }