Skip to content

Commit

Permalink
Add query execution links to exporter results
Browse files Browse the repository at this point in the history
  • Loading branch information
kgopal492 committed Nov 7, 2023
1 parent 67c9de1 commit 9cfb64e
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 8 deletions.
8 changes: 8 additions & 0 deletions querybook/server/lib/export/base_exporter.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
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
Expand Down Expand Up @@ -160,6 +161,13 @@ def _get_statement_execution_download_url(self, statement_execution_id: int):
return reader.get_download_url()
return None

def _get_query_execution_url_by_statement_id(self, statement_execution_id: int):
statement_execution = logic.get_statement_execution_by_id(
statement_execution_id
)
query_execution_id = statement_execution.query_execution_id
return f"{QuerybookSettings.PUBLIC_URL}/query_execution/{query_execution_id}"

def to_dict(self):
return {
"name": self.exporter_name,
Expand Down
20 changes: 17 additions & 3 deletions querybook/server/lib/export/exporters/gspread_exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand All @@ -154,16 +165,19 @@ 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)}",
) as sheet:
self.write_csv_to_sheet(
sheet,
statement_execution_id,
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)
Expand Down
9 changes: 4 additions & 5 deletions querybook/server/lib/export/exporters/python_exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)}
url = "{download_url}"
import requests
import pandas
Expand All @@ -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
)
"""

0 comments on commit 9cfb64e

Please sign in to comment.