From 27c93f438af3538342dc08a6c4905259eb612286 Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Mon, 6 May 2024 09:24:25 -0700
Subject: [PATCH] chore: Add Apache Spark Jinja template processor (#28335)

---
 superset/jinja_context.py                     | 14 ++++++++++++++
 tests/integration_tests/test_jinja_context.py | 17 +++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/superset/jinja_context.py b/superset/jinja_context.py
index 53325b7a4..4a036d769 100644
--- a/superset/jinja_context.py
+++ b/superset/jinja_context.py
@@ -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"
 
@@ -657,6 +670,7 @@ class TrinoTemplateProcessor(PrestoTemplateProcessor):
 DEFAULT_PROCESSORS = {
     "presto": PrestoTemplateProcessor,
     "hive": HiveTemplateProcessor,
+    "spark": SparkTemplateProcessor,
     "trino": TrinoTemplateProcessor,
 }
 
diff --git a/tests/integration_tests/test_jinja_context.py b/tests/integration_tests/test_jinja_context.py
index 8fa61841c..6ebded778 100644
--- a/tests/integration_tests/test_jinja_context.py
+++ b/tests/integration_tests/test_jinja_context.py
@@ -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"