From 9c322ae799dba70b4a28ca38be94f6f0508cf8ff Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <jianchengzhong@pinterest.com>
Date: Thu, 28 Sep 2023 10:29:56 -0700
Subject: [PATCH] chore: add query to ai assistant context logging (#1342)

---
 .../AIAssistant/QueryGenerationModal.tsx          | 15 +++++++++++++++
 querybook/webapp/const/analytics.ts               |  1 +
 2 files changed, 16 insertions(+)

diff --git a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx
index 8136e1b03..bd6b0c3d5 100644
--- a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx
+++ b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx
@@ -240,6 +240,18 @@ export const QueryGenerationModal = ({
             <Button
                 title="Cancel"
                 onClick={() => {
+                    if (newQuery) {
+                        trackClick({
+                            component: ComponentType.AI_ASSISTANT,
+                            element: ElementType.QUERY_GENERATION_REJECT_BUTTON,
+                            aux: {
+                                mode: textToSQLMode,
+                                question,
+                                tables,
+                                query: newQuery,
+                            },
+                        });
+                    }
                     onHide();
                 }}
             />
@@ -255,6 +267,7 @@ export const QueryGenerationModal = ({
                             mode: textToSQLMode,
                             question,
                             tables,
+                            query: newQuery,
                         },
                     });
                     onHide();
@@ -273,6 +286,7 @@ export const QueryGenerationModal = ({
                             mode: textToSQLMode,
                             question,
                             tables,
+                            query: newQuery,
                         },
                     });
                     onHide();
@@ -367,6 +381,7 @@ export const QueryGenerationModal = ({
                                                     mode: textToSQLMode,
                                                     question,
                                                     tables,
+                                                    query: newQuery,
                                                 },
                                             });
                                         }}
diff --git a/querybook/webapp/const/analytics.ts b/querybook/webapp/const/analytics.ts
index 539136eed..370282448 100644
--- a/querybook/webapp/const/analytics.ts
+++ b/querybook/webapp/const/analytics.ts
@@ -103,6 +103,7 @@ export enum ElementType {
     QUERY_ERROR_AUTO_FIX_APPLY_AND_RUN_BUTTON = 'QUERY_ERROR_AUTO_FIX_APPLY_AND_RUN_BUTTON',
     QUERY_GENERATION_BUTTON = 'QUERY_GENERATION_BUTTON',
     QUERY_GENERATION_KEEP_BUTTON = 'QUERY_GENERATION_KEEP_BUTTON',
+    QUERY_GENERATION_REJECT_BUTTON = 'QUERY_GENERATION_REJECT_BUTTON',
     QUERY_GENERATION_APPLY_BUTTON = 'QUERY_GENERATION_APPLY_BUTTON',
     QUERY_GENERATION_APPLY_AND_RUN_BUTTON = 'QUERY_GENERATION_APPLY_AND_RUN_BUTTON',
 }