From 206c856423428519b7b035b2447825cc8d392b48 Mon Sep 17 00:00:00 2001
From: Krishna Gopal <kgopal@pinterest.com>
Date: Wed, 8 Nov 2023 16:09:19 -0800
Subject: [PATCH] Add query execution links to exporter results (#1364)

* Add query execution links to exporter results

* fix pr comments

* update python version in lint workflow
---
 .github/workflows/lint.yml                    |  2 ++
 querybook/server/lib/export/base_exporter.py  | 22 +++++++++++++++++++
 .../lib/export/exporters/gspread_exporter.py  | 20 ++++++++++++++---
 .../lib/export/exporters/python_exporter.py   |  9 ++++----
 4 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 3f5a38840..17318b4a0 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -10,4 +10,6 @@ jobs:
         steps:
             - uses: actions/checkout@v2
             - uses: actions/setup-python@v2
+              with:
+                  python-version: '3.9'
             - uses: pre-commit/action@v2.0.0
diff --git a/querybook/server/lib/export/base_exporter.py b/querybook/server/lib/export/base_exporter.py
index e8db6d704..b5b66f2c9 100644
--- a/querybook/server/lib/export/base_exporter.py
+++ b/querybook/server/lib/export/base_exporter.py
@@ -1,9 +1,13 @@
 from abc import ABCMeta, abstractmethod
 from typing import Generator, List
 from app.db import with_session
+from env import QuerybookSettings
 from lib.logger import get_logger
 from logic import query_execution as logic
 from lib.result_store import GenericReader
+from logic.query_execution_permission import (
+    get_default_user_environment_by_execution_id,
+)
 
 
 LOG = get_logger(__file__)
@@ -160,6 +164,24 @@ def _get_statement_execution_download_url(self, statement_execution_id: int):
                     return reader.get_download_url()
         return None
 
+    @with_session
+    def _get_query_execution_url_by_statement_id(
+        self, statement_execution_id: int, uid: int, session=None
+    ):
+        statement_execution = logic.get_statement_execution_by_id(
+            statement_execution_id, session=session
+        )
+        query_execution_id = statement_execution.query_execution_id
+
+        env = get_default_user_environment_by_execution_id(
+            query_execution_id, uid, session=session
+        )
+        return (
+            f"{QuerybookSettings.PUBLIC_URL}/{env.name}/query_execution/{query_execution_id}/"
+            if env
+            else None
+        )
+
     def to_dict(self):
         return {
             "name": self.exporter_name,
diff --git a/querybook/server/lib/export/exporters/gspread_exporter.py b/querybook/server/lib/export/exporters/gspread_exporter.py
index ce5feb30e..82c881bad 100644
--- a/querybook/server/lib/export/exporters/gspread_exporter.py
+++ b/querybook/server/lib/export/exporters/gspread_exporter.py
@@ -19,6 +19,10 @@
 from logic.user import get_user_by_id, update_user_properties
 from lib.export.base_exporter import BaseExporter
 from lib.form import StructFormField, FormField
+from logic.query_execution import (
+    get_statement_execution_by_id,
+    update_statement_execution,
+)
 
 
 class UserTokenNotFound(Exception):
@@ -141,6 +145,13 @@ def _get_max_rows(
         )
         return max_result_rows - row_offset
 
+    def _save_sheet_to_statement_meta(self, sheet_url: str, statement_id: int):
+        statement_execution = get_statement_execution_by_id(statement_id)
+        meta_info = (
+            statement_execution.meta_info or ""
+        ) + f"Google sheet url: {sheet_url}\n"
+        update_statement_execution(statement_id, meta_info=meta_info)
+
     def export(
         self,
         statement_execution_id,
@@ -154,7 +165,9 @@ def export(
             credentials = self.get_credentials(uid)
             gc = gspread.authorize(credentials)
             with gspread_sheet(
-                gc, sheet_url, f"Querybook Result {statement_execution_id}"
+                gc,
+                sheet_url,
+                f"Querybook Result {statement_execution_id}, {self._get_query_execution_url_by_statement_id(statement_execution_id, uid)}",
             ) as sheet:
                 self.write_csv_to_sheet(
                     sheet,
@@ -162,8 +175,9 @@ def export(
                     worksheet_title,
                     start_cell,
                 )
-
-            return f"https://docs.google.com/spreadsheets/d/{sheet.id}"
+            sheet_url = f"https://docs.google.com/spreadsheets/d/{sheet.id}"
+            self._save_sheet_to_statement_meta(sheet_url, statement_execution_id)
+            return sheet_url
         except RefreshError:
             # Invalidate user access token
             update_user_properties(uid, gspread_token=None)
diff --git a/querybook/server/lib/export/exporters/python_exporter.py b/querybook/server/lib/export/exporters/python_exporter.py
index b25f63cf0..a8a0c9a35 100644
--- a/querybook/server/lib/export/exporters/python_exporter.py
+++ b/querybook/server/lib/export/exporters/python_exporter.py
@@ -15,8 +15,9 @@ def export(self, statement_execution_id, uid):
             statement_execution_id
         )
 
-        return """
-url = "{}"
+        return f"""
+# Querybook execution link: {self._get_query_execution_url_by_statement_id(statement_execution_id, uid)}
+url = "{download_url}"
 
 import requests
 import pandas
@@ -26,6 +27,4 @@ def export(self, statement_execution_id, uid):
     download = s.get(url)
     decoded_content = download.content.decode('utf-8')
     df = pandas.read_csv(StringIO(decoded_content))
-        """.format(
-            download_url
-        )
+        """