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 <2060045+g...@users.noreply.github.com> AuthorDate: Thu Jul 18 07:54:24 2024 -0400 feat: support rendering ObjectStoragePath value (#40638) ] --------- Co-authored-by: Young Li <g...@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()