Skip to content

Commit

Permalink
chore: Add Apache Spark Jinja template processor (#28335)
Browse files Browse the repository at this point in the history
  • Loading branch information
john-bodley authored May 6, 2024
1 parent bf5ff55 commit 27c93f4
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 0 deletions.
14 changes: 14 additions & 0 deletions superset/jinja_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -641,6 +641,19 @@ class HiveTemplateProcessor(PrestoTemplateProcessor):
engine = "hive"


class SparkTemplateProcessor(HiveTemplateProcessor):
engine = "spark"

def process_template(self, sql: str, **kwargs: Any) -> str:
template = self.env.from_string(sql)
kwargs.update(self._context)

# Backwards compatibility if migrating from Hive.
context = validate_template_context(self.engine, kwargs)
context["hive"] = context["spark"]
return template.render(context)


class TrinoTemplateProcessor(PrestoTemplateProcessor):
engine = "trino"

Expand All @@ -657,6 +670,7 @@ def process_template(self, sql: str, **kwargs: Any) -> str:
DEFAULT_PROCESSORS = {
"presto": PrestoTemplateProcessor,
"hive": HiveTemplateProcessor,
"spark": SparkTemplateProcessor,
"trino": TrinoTemplateProcessor,
}

Expand Down
17 changes: 17 additions & 0 deletions tests/integration_tests/test_jinja_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,23 @@ def test_template_hive(app_context: AppContext, mocker: MockFixture) -> None:
assert tp.process_template(template) == "the_latest"


def test_template_spark(app_context: AppContext, mocker: MockFixture) -> None:
lp_mock = mocker.patch(
"superset.jinja_context.SparkTemplateProcessor.latest_partition"
)
lp_mock.return_value = "the_latest"
database = mock.Mock()
database.backend = "spark"
template = "{{ spark.latest_partition('my_table') }}"
tp = get_template_processor(database=database)
assert tp.process_template(template) == "the_latest"

# Backwards compatibility if migrating from 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"
Expand Down

0 comments on commit 27c93f4

Please sign in to comment.