This is an automated email from the ASF dual-hosted git repository.
dpgaspar pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git
The following commit(s) were added to refs/heads/master by this push:
new 72f051f [sqllab] Add CUSTOM_TEMPLATE_PROCESSOR config (#9376)
72f051f is described below
commit 72f051f3ce6b269a0d80de0cef3f357826142213
Author: dandanhub <[email protected]>
AuthorDate: Tue Apr 7 13:00:42 2020 -0700
[sqllab] Add CUSTOM_TEMPLATE_PROCESSOR config (#9376)
Co-authored-by: Dandan Shi <[email protected]>
---
docs/installation.rst | 53 +++++++++++++++
docs/sqllab.rst | 9 +++
superset/config.py | 10 +++
superset/extensions.py | 15 +++-
superset/jinja_context.py | 4 +-
tests/base_tests.py | 22 ++++++
tests/core_tests.py | 83 +++++++++++++++++++++++
tests/superset_test_config.py | 5 ++
tests/superset_test_custom_template_processors.py | 59 ++++++++++++++++
9 files changed, 258 insertions(+), 2 deletions(-)
diff --git a/docs/installation.rst b/docs/installation.rst
index 84dd6e1..3528057 100644
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -1087,6 +1087,59 @@ in this dictionary are made available for users to use
in their SQL.
'my_crazy_macro': lambda x: x*2,
}
+Besides default Jinja templating, SQL lab also supports self-defined template
+processor by setting the ``CUSTOM_TEMPLATE_PROCESSORS`` in your superset
configuration.
+The values in this dictionary overwrite the default Jinja template processors
of the
+specified database engine.
+The example below configures a custom presto template processor which
implements
+its own logic of processing macro template with regex parsing. It uses ``$``
style
+macro instead of ``{{ }}`` style in Jinja templating. By configuring it with
+``CUSTOM_TEMPLATE_PROCESSORS``, sql template on presto database is processed
+by the custom one rather than the default one.
+
+.. code-block:: python
+
+ def DATE(
+ ts: datetime, day_offset: SupportsInt = 0, hour_offset: SupportsInt = 0
+ ) -> str:
+ """Current day as a string."""
+ day_offset, hour_offset = int(day_offset), int(hour_offset)
+ offset_day = (ts + timedelta(days=day_offset,
hours=hour_offset)).date()
+ return str(offset_day)
+
+ class CustomPrestoTemplateProcessor(PrestoTemplateProcessor):
+ """A custom presto template processor."""
+
+ engine = "presto"
+
+ def process_template(self, sql: str, **kwargs) -> str:
+ """Processes a sql template with $ style macro using regex."""
+ # Add custom macros functions.
+ macros = {
+ "DATE": partial(DATE, datetime.utcnow())
+ } # type: Dict[str, Any]
+ # Update with macros defined in context and kwargs.
+ macros.update(self.context)
+ macros.update(kwargs)
+
+ def replacer(match):
+ """Expand $ style macros with corresponding function calls."""
+ macro_name, args_str = match.groups()
+ args = [a.strip() for a in args_str.split(",")]
+ if args == [""]:
+ args = []
+ f = macros[macro_name[1:]]
+ return f(*args)
+
+ macro_names = ["$" + name for name in macros.keys()]
+ pattern = r"(%s)\s*\(([^()]*)\)" % "|".join(map(re.escape,
macro_names))
+ return re.sub(pattern, replacer, sql)
+
+ CUSTOM_TEMPLATE_PROCESSORS = {
+ CustomPrestoTemplateProcessor.engine: CustomPrestoTemplateProcessor
+ }
+
+
SQL Lab also includes a live query validation feature with pluggable backends.
You can configure which validation implementation is used with which database
engine by adding a block like the following to your config.py:
diff --git a/docs/sqllab.rst b/docs/sqllab.rst
index 992a689..aace28f 100644
--- a/docs/sqllab.rst
+++ b/docs/sqllab.rst
@@ -104,6 +104,15 @@ environment using the configuration variable
``JINJA_CONTEXT_ADDONS``.
All objects referenced in this dictionary will become available for users
to integrate in their queries in **SQL Lab**.
+Customize templating
+''''''''''''''''''''
+
+As mentioned in the `Installation & Configuration
<https://superset.incubator.apache.org/installation.html#sql-lab>`__
documentation,
+it's possible for administrators to overwrite Jinja templating with your
customized
+template processor using the configuration variable
``CUSTOM_TEMPLATE_PROCESSORS``.
+The template processors referenced in the dictionary will overwrite default
Jinja template processors
+of the specified database engines.
+
Query cost estimation
'''''''''''''''''''''
diff --git a/superset/config.py b/superset/config.py
index 1cef75f..1d27119 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -34,6 +34,9 @@ from celery.schedules import crontab
from dateutil import tz
from flask_appbuilder.security.manager import AUTH_DB
+from superset.jinja_context import ( # pylint: disable=unused-import
+ BaseTemplateProcessor,
+)
from superset.stats_logger import DummyStatsLogger
from superset.typing import CacheConfig
from superset.utils.log import DBEventLogger
@@ -585,6 +588,13 @@ UPLOADED_CSV_HIVE_NAMESPACE = None
# dictionary.
JINJA_CONTEXT_ADDONS: Dict[str, Callable] = {}
+# A dictionary of macro template processors that gets merged into global
+# template processors. The existing template processors get updated with this
+# dictionary, which means the existing keys get overwritten by the content of
this
+# dictionary. The customized addons don't necessarily need to use jinjia
templating
+# language. This allows you to define custom logic to process macro template.
+CUSTOM_TEMPLATE_PROCESSORS = {} # type: Dict[str, BaseTemplateProcessor]
+
# Roles that are controlled by the API / Superset and should not be changes
# by humans.
ROBOT_PERMISSION_ROLES = ["Public", "Gamma", "Alpha", "Admin", "sql_lab"]
diff --git a/superset/extensions.py b/superset/extensions.py
index 0b7f39b..c501eeb 100644
--- a/superset/extensions.py
+++ b/superset/extensions.py
@@ -20,6 +20,7 @@ import random
import time
import uuid
from datetime import datetime, timedelta
+from typing import Dict, TYPE_CHECKING # pylint: disable=unused-import
import celery
from dateutil.relativedelta import relativedelta
@@ -31,6 +32,12 @@ from werkzeug.local import LocalProxy
from superset.utils.cache_manager import CacheManager
from superset.utils.feature_flag_manager import FeatureFlagManager
+# Avoid circular import
+if TYPE_CHECKING:
+ from superset.jinja_context import ( # pylint: disable=unused-import
+ BaseTemplateProcessor,
+ )
+
class JinjaContextManager:
def __init__(self) -> None:
@@ -42,14 +49,20 @@ class JinjaContextManager:
"timedelta": timedelta,
"uuid": uuid,
}
+ self._template_processors = {} # type: Dict[str,
BaseTemplateProcessor]
def init_app(self, app):
self._base_context.update(app.config["JINJA_CONTEXT_ADDONS"])
+
self._template_processors.update(app.config["CUSTOM_TEMPLATE_PROCESSORS"])
@property
def base_context(self):
return self._base_context
+ @property
+ def template_processors(self):
+ return self._template_processors
+
class ResultsBackendManager:
def __init__(self) -> None:
@@ -120,7 +133,7 @@ db = SQLA()
_event_logger: dict = {}
event_logger = LocalProxy(lambda: _event_logger.get("event_logger"))
feature_flag_manager = FeatureFlagManager()
-jinja_context_manager = JinjaContextManager()
+jinja_context_manager = JinjaContextManager() # type: JinjaContextManager
manifest_processor = UIManifestProcessor(APP_DIR)
migrate = Migrate()
results_backend_manager = ResultsBackendManager()
diff --git a/superset/jinja_context.py b/superset/jinja_context.py
index ebdd665..7863809 100644
--- a/superset/jinja_context.py
+++ b/superset/jinja_context.py
@@ -23,6 +23,7 @@ from flask import g, request
from jinja2.sandbox import SandboxedEnvironment
from superset import jinja_base_context
+from superset.extensions import jinja_context_manager
def url_param(param: str, default: Optional[str] = None) -> Optional[Any]:
@@ -263,7 +264,8 @@ class HiveTemplateProcessor(PrestoTemplateProcessor):
engine = "hive"
-template_processors = {}
+# The global template processors from Jinja context manager.
+template_processors = jinja_context_manager.template_processors
keys = tuple(globals().keys())
for k in keys:
o = globals()[k]
diff --git a/tests/base_tests.py b/tests/base_tests.py
index 7ccbf0a..97c69f3 100644
--- a/tests/base_tests.py
+++ b/tests/base_tests.py
@@ -282,6 +282,28 @@ class SupersetTestCase(TestCase):
if database:
db.session.delete(database)
+ def create_fake_presto_db(self):
+ self.login(username="admin")
+ database_name = "presto"
+ db_id = 200
+ return self.get_or_create(
+ cls=models.Database,
+ criteria={"database_name": database_name},
+ session=db.session,
+ sqlalchemy_uri="presto://user@host:8080/hive",
+ id=db_id,
+ )
+
+ def delete_fake_presto_db(self):
+ database = (
+ db.session.query(Database)
+ .filter(Database.database_name == "presto")
+ .scalar()
+ )
+ if database:
+ db.session.delete(database)
+ db.session.commit()
+
def validate_sql(
self,
sql,
diff --git a/tests/core_tests.py b/tests/core_tests.py
index 9c6b54a..eb3e2f7 100644
--- a/tests/core_tests.py
+++ b/tests/core_tests.py
@@ -668,6 +668,89 @@ class CoreTests(SupersetTestCase):
data = self.run_sql(sql, "fdaklj3ws")
self.assertEqual(data["data"][0]["test"], "2017-01-01T00:00:00")
+ @mock.patch("tests.superset_test_custom_template_processors.datetime")
+ def test_custom_process_template(self, mock_dt) -> None:
+ """Test macro defined in custom template processor works."""
+ mock_dt.utcnow = mock.Mock(return_value=datetime.datetime(1970, 1, 1))
+ db = mock.Mock()
+ db.backend = "presto"
+ tp = jinja_context.get_template_processor(database=db)
+
+ sql = "SELECT '$DATE()'"
+ rendered = tp.process_template(sql)
+ self.assertEqual("SELECT '{}'".format("1970-01-01"), rendered)
+
+ sql = "SELECT '$DATE(1, 2)'"
+ rendered = tp.process_template(sql)
+ self.assertEqual("SELECT '{}'".format("1970-01-02"), rendered)
+
+ def test_custom_get_template_kwarg(self):
+ """Test macro passed as kwargs when getting template processor
+ works in custom template processor."""
+ db = mock.Mock()
+ db.backend = "presto"
+ s = "$foo()"
+ tp = jinja_context.get_template_processor(database=db, foo=lambda:
"bar")
+ rendered = tp.process_template(s)
+ self.assertEqual("bar", rendered)
+
+ def test_custom_template_kwarg(self) -> None:
+ """Test macro passed as kwargs when processing template
+ works in custom template processor."""
+ db = mock.Mock()
+ db.backend = "presto"
+ s = "$foo()"
+ tp = jinja_context.get_template_processor(database=db)
+ rendered = tp.process_template(s, foo=lambda: "bar")
+ self.assertEqual("bar", rendered)
+
+ def test_custom_template_processors_overwrite(self) -> None:
+ """Test template processor for presto gets overwritten by custom
one."""
+ db = mock.Mock()
+ db.backend = "presto"
+ tp = jinja_context.get_template_processor(database=db)
+
+ sql = "SELECT '{{ datetime(2017, 1, 1).isoformat() }}'"
+ rendered = tp.process_template(sql)
+ self.assertEqual(sql, rendered)
+
+ sql = "SELECT '{{ DATE(1, 2) }}'"
+ rendered = tp.process_template(sql)
+ self.assertEqual(sql, rendered)
+
+ def test_custom_template_processors_ignored(self) -> None:
+ """Test custom template processor is ignored for a difference backend
+ database."""
+ maindb = utils.get_example_database()
+ sql = "SELECT '$DATE()'"
+ tp = jinja_context.get_template_processor(database=maindb)
+ rendered = tp.process_template(sql)
+ self.assertEqual(sql, rendered)
+
+ @mock.patch("tests.superset_test_custom_template_processors.datetime")
+ @mock.patch("superset.sql_lab.get_sql_results")
+ def test_custom_templated_sql_json(self, sql_lab_mock, mock_dt) -> None:
+ """Test sqllab receives macros expanded query."""
+ mock_dt.utcnow = mock.Mock(return_value=datetime.datetime(1970, 1, 1))
+ self.login("admin")
+ sql = "SELECT '$DATE()' as test"
+ resp = {
+ "status": utils.QueryStatus.SUCCESS,
+ "query": {"rows": 1},
+ "data": [{"test": "'1970-01-01'"}],
+ }
+ sql_lab_mock.return_value = resp
+
+ dbobj = self.create_fake_presto_db()
+ json_payload = dict(database_id=dbobj.id, sql=sql)
+ self.get_json_resp(
+ "/superset/sql_json/", raise_on_error=False, json_=json_payload
+ )
+ assert sql_lab_mock.called
+ self.assertEqual(sql_lab_mock.call_args[0][1], "SELECT '1970-01-01' as
test")
+
+ self.delete_fake_presto_db()
+
def test_fetch_datasource_metadata(self):
self.login(username="admin")
url = "/superset/fetch_datasource_metadata?" "datasourceKey=1__table"
diff --git a/tests/superset_test_config.py b/tests/superset_test_config.py
index ff0a1e5..12d7a06 100644
--- a/tests/superset_test_config.py
+++ b/tests/superset_test_config.py
@@ -18,6 +18,7 @@
from copy import copy
from superset.config import *
+from tests.superset_test_custom_template_processors import
CustomPrestoTemplateProcessor
AUTH_USER_REGISTRATION_ROLE = "alpha"
SQLALCHEMY_DATABASE_URI = "sqlite:///" + os.path.join(DATA_DIR, "unittests.db")
@@ -57,3 +58,7 @@ class CeleryConfig(object):
CELERY_CONFIG = CeleryConfig
+
+CUSTOM_TEMPLATE_PROCESSORS = {
+ CustomPrestoTemplateProcessor.engine: CustomPrestoTemplateProcessor
+}
diff --git a/tests/superset_test_custom_template_processors.py
b/tests/superset_test_custom_template_processors.py
new file mode 100644
index 0000000..28fc65d
--- /dev/null
+++ b/tests/superset_test_custom_template_processors.py
@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import re
+from datetime import datetime, timedelta
+from functools import partial
+from typing import Any, Dict, SupportsInt
+
+from superset.jinja_context import PrestoTemplateProcessor
+
+
+def DATE(
+ ts: datetime, day_offset: SupportsInt = 0, hour_offset: SupportsInt = 0
+) -> str:
+ """Current day as a string"""
+ day_offset, hour_offset = int(day_offset), int(hour_offset)
+ offset_day = (ts + timedelta(days=day_offset, hours=hour_offset)).date()
+ return str(offset_day)
+
+
+class CustomPrestoTemplateProcessor(PrestoTemplateProcessor):
+ """A custom presto template processor for test."""
+
+ engine = "presto"
+
+ def process_template(self, sql: str, **kwargs) -> str:
+ """Processes a sql template with $ style macro using regex."""
+ # Add custom macros functions.
+ macros = {"DATE": partial(DATE, datetime.utcnow())} # type: Dict[str,
Any]
+ # Update with macros defined in context and kwargs.
+ macros.update(self.context)
+ macros.update(kwargs)
+
+ def replacer(match):
+ """Expands $ style macros with corresponding function calls."""
+ macro_name, args_str = match.groups()
+ args = [a.strip() for a in args_str.split(",")]
+ if args == [""]:
+ args = []
+ f = macros[macro_name[1:]]
+ return f(*args)
+
+ macro_names = ["$" + name for name in macros.keys()]
+ pattern = r"(%s)\s*\(([^()]*)\)" % "|".join(map(re.escape,
macro_names))
+ return re.sub(pattern, replacer, sql)
