From b8d9a95e0636421459d50da44c73a176a5e50c63 Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <jianchengzhong@pinterest.com>
Date: Thu, 16 May 2024 11:05:38 -0700
Subject: [PATCH] fix: add back text2sql survey (#1452)

* fix: add back text2sql survey

* fix linter
---
 package.json                                      |  2 +-
 .../AIAssistant/AICommandResultView.tsx           | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/package.json b/package.json
index daf51eff0..bf63f6f3e 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.34.0",
+    "version": "3.34.1",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {
diff --git a/querybook/webapp/components/AIAssistant/AICommandResultView.tsx b/querybook/webapp/components/AIAssistant/AICommandResultView.tsx
index 8effd03d1..dbfe63025 100644
--- a/querybook/webapp/components/AIAssistant/AICommandResultView.tsx
+++ b/querybook/webapp/components/AIAssistant/AICommandResultView.tsx
@@ -7,6 +7,8 @@ import { IQueryCellCommand } from 'const/command';
 import useNonEmptyState from 'hooks/useNonEmptyState';
 import { trackClick } from 'lib/analytics';
 import { Button } from 'ui/Button/Button';
+import { SurveySurfaceType } from 'const/survey';
+import { useSurveyTrigger } from 'hooks/ui/useSurveyTrigger';
 
 interface IAICommandResultViewProps {
     command: IQueryCellCommand;
@@ -59,6 +61,19 @@ export const AICommandResultView = ({
         }
     }, [commandResult]);
 
+    const triggerSurvey = useSurveyTrigger();
+    useEffect(() => {
+        if (!newQuery || isStreaming) {
+            return;
+        }
+        triggerSurvey(SurveySurfaceType.TEXT_TO_SQL, {
+            question: commandKwargs.question,
+            tables,
+            query: newQuery,
+        });
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [newQuery, triggerSurvey, isStreaming]);
+
     const handleAccept = useCallback(() => {
         onAccept(newQuery);
         trackClick({