feat: add drag and drop for templated variables

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.15.4",
+    "version": "3.15.5",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {
@@ -54,7 +54,7 @@
         "draft-js-export-html": "^1.4.1",
         "draft-js-import-html": "^1.4.1",
         "fast-json-stable-stringify": "2.0.0",
-        "formik": "2.2.6",
+        "formik": "2.2.9",
         "history": "^4.10.1",
         "html-webpack-plugin": "5.3.1",
         "immer": "8.0.1",

querybook/server/datasources/query_execution.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 from flask import abort, Response, redirect
 from flask_login import current_user
 
@@ -20,6 +18,8 @@
     render_templated_query,
 )
 from lib.form import validate_form
+from lib.data_doc.meta import var_config_to_var_dict
+from lib.data_doc.doc_types import DataDocMetaVarConfig
 from const.query_execution import QueryExecutionExportStatus, QueryExecutionStatus
 from const.datasources import RESOURCE_NOT_FOUND_STATUS_CODE
 from logic import (
@@ -450,9 +450,13 @@ def poll_export_statement_execution_result(task_id):
 
 
 @register("/query_execution/templated_query/", methods=["POST"])
-def get_templated_query(query: str, variables: Dict[str, str], engine_id: int):
+def get_templated_query(
+    query: str, var_config: list[DataDocMetaVarConfig], engine_id: int
+):
     try:
-        return render_templated_query(query, variables, engine_id)
+        return render_templated_query(
+            query, var_config_to_var_dict(var_config), engine_id
+        )
     except QueryTemplatingError as e:
         raise RequestException(e)
 

querybook/server/lib/dag_exporter/export_dag.py

@@ -21,7 +21,7 @@ def export_dag(data_doc_id, dag_exporter_name, session=None):
         edges=dag["edges"],
         meta={
             **dag_export["meta"]["exporter_meta"][dag_exporter_name],
-            "variables": doc.meta,
+            "variables": doc.meta_variables,
         },
         cell_by_id=cell_by_id,
     )

querybook/server/lib/data_doc/doc_types.py

@@ -0,0 +1,11 @@
+from typing import TypedDict, Union, Literal
+
+
+class DataDocMetaVarConfig(TypedDict):
+    name: str
+    value: Union[str, int, float, bool]
+    type: Literal["boolean", "number", "string"]
+
+
+class DataDocMeta(TypedDict):
+    variables: list[DataDocMetaVarConfig]

querybook/server/lib/data_doc/meta.py

@@ -0,0 +1,83 @@
+from typing import Dict, Any
+from .doc_types import DataDocMeta, DataDocMetaVarConfig
+
+
+def check_variable_type(val: Any):
+    if isinstance(val, (int, float)):
+        return "number"
+    elif isinstance(val, bool):
+        return "boolean"
+    elif isinstance(val, str):
+        return "string"
+
+    # this shouldn't happen, just in case
+    return "string"
+
+
+def convert_if_legacy_datadoc_meta_v0(datadoc_meta: Dict) -> DataDocMeta:
+    """Converts the old meta format which is only a dictionary of templated variables
+       to a more general format that has templated vars as array plus other fields
+
+        Old meta: `{ "foo": "bar" }`
+        New meta: `{ "variables": ["name": "foo", "type": "string", "value": "bar", ...] }`
+
+        If the new meta is passed in, no change would be made.
+
+    Args:
+        datadoc_meta (Dict): Old/New meta format
+
+    Returns:
+        DataDocMeta: New meta format
+    """
+    if isinstance(datadoc_meta.get("variables"), list):
+        return datadoc_meta
+
+    new_meta = {"variables": []}
+
+    for name, value in datadoc_meta.items():
+        new_meta["variables"].append(
+            {"name": name, "value": value, "type": check_variable_type(value)}
+        )
+
+    return new_meta
+
+
+def convert_if_legacy_datadoc_meta(datadoc_meta: Dict) -> DataDocMeta:
+    datadoc_meta = convert_if_legacy_datadoc_meta_v0(datadoc_meta)
+    return datadoc_meta
+
+
+def var_config_to_var_dict(variables: list[DataDocMetaVarConfig]) -> Dict:
+    var_dict = {}
+
+    for config in variables:
+        var_dict[config["name"]] = config["value"]
+
+    return var_dict
+
+
+valid_meta_keys = ["variables"]
+
+
+def validate_datadoc_meta(datadoc_meta: DataDocMeta) -> bool:
+    for key in datadoc_meta.keys():
+        if key not in valid_meta_keys:
+            return False
+
+    if "variables" in datadoc_meta:
+        variables = datadoc_meta["variables"]
+        if not isinstance(variables, list):
+            return False
+
+        for variable_config in variables:
+            var_type = variable_config["type"]
+            var_val = variable_config["value"]
+
+            if var_type == "string" and not isinstance(var_val, str):
+                return False
+            if var_type == "boolean" and not isinstance(var_val, bool):
+                return False
+            if var_type == "number" and not isinstance(var_val, (float, int)):
+                return False
+
+    return True

querybook/server/logic/datadoc.py

@@ -49,16 +49,18 @@ def create_data_doc(
     commit=True,
     session=None,
 ):
-    data_doc = DataDoc(
-        public=public,
-        archived=archived,
-        owner_uid=owner_uid,
-        environment_id=environment_id,
-        title=title,
-        meta=meta,
+    data_doc = DataDoc.create(
+        fields={
+            "public": public,
+            "archived": archived,
+            "owner_uid": owner_uid,
+            "environment_id": environment_id,
+            "title": title,
+            "meta": meta,
+        },
+        commit=False,
+        session=session,
     )
-    session.add(data_doc)
-    session.flush()
 
     for index, cell in enumerate(cells):
         data_cell = create_data_cell(
@@ -150,7 +152,12 @@ def update_data_doc(id, commit=True, session=None, **fields):
 
         if commit:
             session.commit()
-            update_es_data_doc_by_id(data_doc.id)
+
+            if any(
+                field_name in fields
+                for field_name in ["public", "archived", "owner_uid", "title"]
+            ):
+                update_es_data_doc_by_id(data_doc.id)
 
             # update es queries if doc is switched between public/private
             if "public" in fields:

querybook/server/models/datadoc.py

@@ -1,10 +1,17 @@
 import sqlalchemy as sql
 from sqlalchemy.orm import backref, relationship
+from sqlalchemy.ext.hybrid import hybrid_property
 
 from app import db
 from const.db import name_length, now, description_length, mediumtext_length
 from const.data_doc import DataCellType
 from lib.sqlalchemy import CRUDMixin
+from lib.data_doc.meta import (
+    convert_if_legacy_datadoc_meta,
+    var_config_to_var_dict,
+    validate_datadoc_meta,
+)
+from lib.data_doc.doc_types import DataDocMeta
 
 Base = db.Base
 
@@ -31,7 +38,7 @@ class DataDoc(Base, CRUDMixin):
     updated_at = sql.Column(sql.DateTime, default=now, nullable=False)
 
     title = sql.Column(sql.String(length=name_length), default="", nullable=False)
-    meta = sql.Column(sql.JSON, default={}, nullable=False)
+    _meta = sql.Column("meta", sql.JSON, default={}, nullable=False)
 
     cells = relationship(
         "DataCell",
@@ -48,6 +55,26 @@ class DataDoc(Base, CRUDMixin):
         backref=backref("data_docs", cascade="all, delete", passive_deletes=True),
     )
 
+    @hybrid_property
+    def meta(self) -> DataDocMeta:
+        return convert_if_legacy_datadoc_meta(self._meta or {})
+
+    @meta.setter
+    def meta(self, new_meta: DataDocMeta):
+        is_valid = validate_datadoc_meta(new_meta)
+        if not is_valid:
+            raise ValueError("Invalid DataDoc.meta")
+
+        self._meta = new_meta
+
+    @property
+    def meta_variables(self) -> dict:
+        """
+        The field is used to generate a dictionary of templated variables.
+        It is used in scheduled data docs
+        """
+        return var_config_to_var_dict(self.meta.get("variables", []))
+
     def to_dict(self, with_cells=False):
         data_doc_dict = {
             "id": self.id,

querybook/server/tasks/run_datadoc.py

@@ -73,7 +73,10 @@ def run_datadoc_with_config(
 
             try:
                 query = render_templated_query(
-                    query_cell.context, data_doc.meta, engine_id, session=session
+                    query_cell.context,
+                    data_doc.meta_variables,
+                    engine_id,
+                    session=session,
                 )
             except Exception as e:
                 on_datadoc_completion(

querybook/webapp/components/DataDoc/DataDoc.tsx

@@ -22,6 +22,7 @@ import {
     IDataCell,
     IDataCellMeta,
     IDataDoc,
+    IDataDocMeta,
     IDataQueryCell,
 } from 'const/datadoc';
 import { ISearchOptions, ISearchResult } from 'const/searchAndReplace';
@@ -632,7 +633,7 @@ class DataDocComponent extends React.PureComponent<IProps, IState> {
                     key={cell.id}
                     docId={dataDoc.id}
                     numberOfCells={dataDoc.dataDocCells.length}
-                    templatedVariables={dataDoc.meta}
+                    templatedVariables={dataDoc.meta.variables}
                     cell={cell}
                     index={index}
                     queryIndexInDoc={queryIndexInDoc}
@@ -907,7 +908,7 @@ function mapDispatchToProps(dispatch: Dispatch) {
             );
         },
 
-        changeDataDocMeta: (docId: number, meta: IDataCellMeta) =>
+        changeDataDocMeta: (docId: number, meta: IDataDocMeta) =>
             dispatch(dataDocActions.updateDataDocField(docId, 'meta', meta)),
 
         cloneDataDoc: (docId: number) =>

querybook/webapp/components/DataDocCell/DataDocCell.tsx

@@ -8,7 +8,11 @@ import { DataDocChartCell } from 'components/DataDocChartCell/DataDocChartCell';
 import { DataDocQueryCell } from 'components/DataDocQueryCell/DataDocQueryCell';
 import { DataDocTextCell } from 'components/DataDocTextCell/DataDocTextCell';
 import { UserAvatar } from 'components/UserBadge/UserAvatar';
-import { DataCellUpdateFields, IDataCell } from 'const/datadoc';
+import {
+    DataCellUpdateFields,
+    IDataCell,
+    TDataDocMetaVariables,
+} from 'const/datadoc';
 import { DataDocContext } from 'context/DataDoc';
 import { useMakeSelector } from 'hooks/redux/useMakeSelector';
 import { useBoundFunc } from 'hooks/useBoundFunction';
@@ -22,7 +26,7 @@ import './DataDocCell.scss';
 interface IDataDocCellProps {
     docId: number;
     numberOfCells: number;
-    templatedVariables: Record<string, string>;
+    templatedVariables: TDataDocMetaVariables;
 
     cell: IDataCell;
     index: number;

querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx

@@ -21,7 +21,7 @@ import { QuerySnippetInsertionModal } from 'components/QuerySnippetInsertionModa
 import { TemplatedQueryView } from 'components/TemplateQueryView/TemplatedQueryView';
 import { TranspileQueryModal } from 'components/TranspileQueryModal/TranspileQueryModal';
 import { UDFForm } from 'components/UDFForm/UDFForm';
-import { IDataQueryCellMeta } from 'const/datadoc';
+import { IDataQueryCellMeta, TDataDocMetaVariables } from 'const/datadoc';
 import type { IQueryEngine, IQueryTranspiler } from 'const/queryEngine';
 import CodeMirror from 'lib/codemirror';
 import { createSQLLinter } from 'lib/codemirror/codemirror-lint';
@@ -73,7 +73,7 @@ interface IOwnProps {
     cellId: number;
 
     queryIndexInDoc: number;
-    templatedVariables: Record<string, string>;
+    templatedVariables: TDataDocMetaVariables;
 
     shouldFocus: boolean;
     isFullScreen: boolean;
@@ -364,7 +364,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
 
     @bind
     public async getTransformedQuery() {
-        const { templatedVariables = {} } = this.props;
+        const { templatedVariables = [] } = this.props;
         const { query } = this.state;
         const selectedRange =
             this.queryEditorRef.current &&

querybook/webapp/components/DataDocRightSidebar/DataDocRightSidebar.tsx

@@ -5,7 +5,7 @@ import { DataDocBoardsButton } from 'components/DataDocBoardsButton/DataDocBoard
 import { DataDocDAGExporterButton } from 'components/DataDocDAGExporter/DataDocDAGExporterButton';
 import { DataDocTemplateButton } from 'components/DataDocTemplateButton/DataDocTemplateButton';
 import { DataDocUIGuide } from 'components/UIGuide/DataDocUIGuide';
-import { IDataDoc } from 'const/datadoc';
+import { IDataDoc, IDataDocMeta } from 'const/datadoc';
 import { useAnnouncements } from 'hooks/redux/useAnnouncements';
 import { useScrollToTop } from 'hooks/ui/useScrollToTop';
 import { fetchDAGExporters } from 'redux/dataDoc/action';
@@ -24,7 +24,7 @@ interface IProps {
     isEditable: boolean;
     isConnected: boolean;
 
-    changeDataDocMeta: (docId: number, meta: Record<string, any>) => any;
+    changeDataDocMeta: (docId: number, meta: IDataDocMeta) => Promise<void>;
     onClone: () => any;
 
     onCollapse: () => any;

querybook/webapp/components/DataDocTemplateButton/DataDocTemplateButton.tsx

@@ -1,15 +1,14 @@
 import React from 'react';
-import toast from 'react-hot-toast';
 
 import { DataDocTemplateVarForm } from 'components/DataDocTemplateButton/DataDocTemplateVarForm';
-import { IDataDoc } from 'const/datadoc';
+import { IDataDoc, IDataDocMeta } from 'const/datadoc';
 import { IconButton } from 'ui/Button/IconButton';
 import { Modal } from 'ui/Modal/Modal';
 
 import { DataDocTemplateInfoButton } from './DataDocTemplateInfoButton';
 
 interface IProps {
-    changeDataDocMeta: (docId: number, meta: Record<string, any>) => any;
+    changeDataDocMeta: (docId: number, meta: IDataDocMeta) => Promise<void>;
     dataDoc: IDataDoc;
     isEditable?: boolean;
 }
@@ -31,11 +30,13 @@ export const DataDocTemplateButton: React.FunctionComponent<IProps> = ({
         >
             <DataDocTemplateVarForm
                 isEditable={isEditable}
-                templatedVariables={dataDoc.meta}
-                onSave={(meta) => {
-                    changeDataDocMeta(dataDoc.id, meta);
+                variables={dataDoc.meta.variables}
+                onSave={(variables) => {
                     setShowTemplateForm(false);
-                    toast.success('Variables saved');
+                    return changeDataDocMeta(dataDoc.id, {
+                        ...dataDoc.meta,
+                        variables,
+                    });
                 }}
             />
         </Modal>