superset/tests/integration_tests/test_jinja_context.py

208 lines
8.1 KiB
Python

# 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.
from datetime import datetime
from unittest import mock
import pytest
from flask.ctx import AppContext
from pytest_mock import MockFixture
import superset.utils.database
from superset.exceptions import SupersetTemplateException
from superset.jinja_context import get_template_processor
def test_process_template(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "SELECT '{{ 1+1 }}'"
tp = get_template_processor(database=maindb)
assert tp.process_template(template) == "SELECT '2'"
def test_get_template_kwarg(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo }}"
tp = get_template_processor(database=maindb, foo="bar")
assert tp.process_template(template) == "bar"
def test_template_kwarg(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo }}"
tp = get_template_processor(database=maindb)
assert tp.process_template(template, foo="bar") == "bar"
def test_get_template_kwarg_dict(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo.bar }}"
tp = get_template_processor(database=maindb, foo={"bar": "baz"})
assert tp.process_template(template) == "baz"
def test_template_kwarg_dict(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo.bar }}"
tp = get_template_processor(database=maindb)
assert tp.process_template(template, foo={"bar": "baz"}) == "baz"
def test_get_template_kwarg_lambda(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo() }}"
tp = get_template_processor(database=maindb, foo=lambda: "bar")
with pytest.raises(SupersetTemplateException):
tp.process_template(template)
def test_template_kwarg_lambda(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo() }}"
tp = get_template_processor(database=maindb)
with pytest.raises(SupersetTemplateException):
tp.process_template(template, foo=lambda: "bar")
def test_get_template_kwarg_module(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ dt(2017, 1, 1).isoformat() }}"
tp = get_template_processor(database=maindb, dt=datetime)
with pytest.raises(SupersetTemplateException):
tp.process_template(template)
def test_template_kwarg_module(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ dt(2017, 1, 1).isoformat() }}"
tp = get_template_processor(database=maindb)
with pytest.raises(SupersetTemplateException):
tp.process_template(template, dt=datetime)
def test_get_template_kwarg_nested_module(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo.dt }}"
tp = get_template_processor(database=maindb, foo={"dt": datetime})
with pytest.raises(SupersetTemplateException):
tp.process_template(template)
def test_template_kwarg_nested_module(app_context: AppContext) -> None:
maindb = superset.utils.database.get_example_database()
template = "{{ foo.dt }}"
tp = get_template_processor(database=maindb)
with pytest.raises(SupersetTemplateException):
tp.process_template(template, foo={"bar": datetime})
def test_template_hive(app_context: AppContext, mocker: MockFixture) -> None:
lp_mock = mocker.patch(
"superset.jinja_context.HiveTemplateProcessor.latest_partition"
)
lp_mock.return_value = "the_latest"
database = mock.Mock()
database.backend = "hive"
template = "{{ hive.latest_partition('my_table') }}"
tp = get_template_processor(database=database)
assert tp.process_template(template) == "the_latest"
def test_template_trino(app_context: AppContext, mocker: MockFixture) -> None:
lp_mock = mocker.patch(
"superset.jinja_context.TrinoTemplateProcessor.latest_partition"
)
lp_mock.return_value = "the_latest"
database = mock.Mock()
database.backend = "trino"
template = "{{ trino.latest_partition('my_table') }}"
tp = get_template_processor(database=database)
assert tp.process_template(template) == "the_latest"
# Backwards compatibility if migrating from Presto.
template = "{{ presto.latest_partition('my_table') }}"
tp = get_template_processor(database=database)
assert tp.process_template(template) == "the_latest"
def test_template_context_addons(app_context: AppContext, mocker: MockFixture) -> None:
addons_mock = mocker.patch("superset.jinja_context.context_addons")
addons_mock.return_value = {"datetime": datetime}
maindb = superset.utils.database.get_example_database()
template = "SELECT '{{ datetime(2017, 1, 1).isoformat() }}'"
tp = get_template_processor(database=maindb)
assert tp.process_template(template) == "SELECT '2017-01-01T00:00:00'"
def test_custom_process_template(app_context: AppContext, mocker: MockFixture) -> None:
"""Test macro defined in custom template processor works."""
mock_dt = mocker.patch(
"tests.integration_tests.superset_test_custom_template_processors.datetime"
)
mock_dt.utcnow = mock.Mock(return_value=datetime(1970, 1, 1))
database = mock.Mock()
database.backend = "db_for_macros_testing"
tp = get_template_processor(database=database)
template = "SELECT '$DATE()'"
assert tp.process_template(template) == f"SELECT '1970-01-01'"
template = "SELECT '$DATE(1, 2)'"
assert tp.process_template(template) == "SELECT '1970-01-02'"
def test_custom_get_template_kwarg(app_context: AppContext) -> None:
"""Test macro passed as kwargs when getting template processor
works in custom template processor."""
database = mock.Mock()
database.backend = "db_for_macros_testing"
template = "$foo()"
tp = get_template_processor(database=database, foo=lambda: "bar")
assert tp.process_template(template) == "bar"
def test_custom_template_kwarg(app_context: AppContext) -> None:
"""Test macro passed as kwargs when processing template
works in custom template processor."""
database = mock.Mock()
database.backend = "db_for_macros_testing"
template = "$foo()"
tp = get_template_processor(database=database)
assert tp.process_template(template, foo=lambda: "bar") == "bar"
def test_custom_template_processors_overwrite(app_context: AppContext) -> None:
"""Test template processor for presto gets overwritten by custom one."""
database = mock.Mock()
database.backend = "db_for_macros_testing"
tp = get_template_processor(database=database)
template = "SELECT '{{ datetime(2017, 1, 1).isoformat() }}'"
assert tp.process_template(template) == template
template = "SELECT '{{ DATE(1, 2) }}'"
assert tp.process_template(template) == template
def test_custom_template_processors_ignored(app_context: AppContext) -> None:
"""Test custom template processor is ignored for a difference backend
database."""
maindb = superset.utils.database.get_example_database()
template = "SELECT '$DATE()'"
tp = get_template_processor(database=maindb)
assert tp.process_template(template) == template