Skip to content

Commit

Permalink
Add query execution links to exporter results (#1364)
Browse files Browse the repository at this point in the history
* Add query execution links to exporter results

* fix pr comments

* update python version in lint workflow
  • Loading branch information
kgopal492 authored Nov 9, 2023
1 parent 67c9de1 commit 206c856
Show file tree
Hide file tree
Showing 4 changed files with 45 additions and 8 deletions.
2 changes: 2 additions & 0 deletions .github/workflows/lint.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,6 @@ jobs:
steps:
- uses: actions/checkout@v2
- uses: actions/setup-python@v2
with:
python-version: '3.9'
- uses: pre-commit/[email protected]
22 changes: 22 additions & 0 deletions querybook/server/lib/export/base_exporter.py
Original file line number Diff line number Diff line change
@@ -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__)
Expand Down Expand Up @@ -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,
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, uid)}",
) 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, uid)}
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 206c856

Please sign in to comment.