(airflow) branch main updated: feat: support rendering ObjectStoragePath value (#40638)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new f88fdbabb1 feat: support rendering ObjectStoragePath value (#40638)
f88fdbabb1 is described below

commit f88fdbabb10c8661552b2b4fee83d5efd3820326
Author: Guangyang Li <20...@users.noreply.github.com>
AuthorDate: Thu Jul 18 07:54:24 2024 -0400

    feat: support rendering ObjectStoragePath value (#40638)
    
    ]
    
    ---------
    
    Co-authored-by: Young Li <gy...@uber.com>
---
 airflow/template/templater.py    | 11 +++++++++++
 tests/template/test_templater.py |  9 +++++++++
 2 files changed, 20 insertions(+)

diff --git a/airflow/template/templater.py b/airflow/template/templater.py
index 93ca679982..bfbb6f5722 100644
--- a/airflow/template/templater.py
+++ b/airflow/template/templater.py
@@ -20,6 +20,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Collection, Iterable, Sequence
 
+from airflow.io.path import ObjectStoragePath
 from airflow.utils.helpers import render_template_as_native, render_template_to_string
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.mixins import ResolveMixin
@@ -168,6 +169,8 @@ class Templater(LoggingMixin):
             else:
                 template = jinja_env.from_string(value)
             return self._render(template, context)
+        if isinstance(value, ObjectStoragePath):
+            return self._render_object_storage_path(value, context, jinja_env)
         if isinstance(value, ResolveMixin):
             return value.resolve(context)
 
@@ -187,6 +190,14 @@ class Templater(LoggingMixin):
         self._render_nested_template_fields(value, context, jinja_env, oids)
         return value
 
+    def _render_object_storage_path(
+        self, value: ObjectStoragePath, context: Context, jinja_env: jinja2.Environment
+    ) -> ObjectStoragePath:
+        serialized_path = value.serialize()
+        path_version = value.__version__
+        serialized_path["path"] = self._render(jinja_env.from_string(serialized_path["path"]), context)
+        return ObjectStoragePath.deserialize(data=serialized_path, version=path_version)
+
     def _render_nested_template_fields(
         self,
         value: Any,
diff --git a/tests/template/test_templater.py b/tests/template/test_templater.py
index c28e1f0503..e1dd9bedb0 100644
--- a/tests/template/test_templater.py
+++ b/tests/template/test_templater.py
@@ -19,6 +19,7 @@ from __future__ import annotations
 
 import jinja2
 
+from airflow.io.path import ObjectStoragePath
 from airflow.models.dag import DAG
 from airflow.template.templater import LiteralValue, Templater
 from airflow.utils.context import Context
@@ -52,6 +53,14 @@ class TestTemplater:
         templater.resolve_template_files()
         assert "Failed to resolve template field 'message'" in caplog.text
 
+    def test_render_object_storage_path(self):
+        templater = Templater()
+        path = ObjectStoragePath("s3://bucket/key/{{ ds }}/part")
+        context = Context({"ds": "2006-02-01"})  # type: ignore
+        jinja_env = templater.get_template_env()
+        rendered_content = templater._render_object_storage_path(path, context, jinja_env)
+        assert rendered_content == ObjectStoragePath("s3://bucket/key/2006-02-01/part")
+
     def test_render_template(self):
         context = Context({"name": "world"})  # type: ignore
         templater = Templater()