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()
