1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18from __future__ import annotations
19
20from typing import TYPE_CHECKING
21
22import jinja2.nativetypes
23import jinja2.sandbox
24
25if TYPE_CHECKING:
26 import datetime
27
28
29class _AirflowEnvironmentMixin:
30 def __init__(self, **kwargs):
31 super().__init__(**kwargs)
32
33 self.filters.update(FILTERS)
34
35 def is_safe_attribute(self, obj, attr, value):
36 """
37 Allow access to ``_`` prefix vars (but not ``__``).
38
39 Unlike the stock SandboxedEnvironment, we allow access to "private" attributes (ones starting with
40 ``_``) whilst still blocking internal or truly private attributes (``__`` prefixed ones).
41 """
42 return not jinja2.sandbox.is_internal_attribute(obj, attr)
43
44
45class NativeEnvironment(_AirflowEnvironmentMixin, jinja2.nativetypes.NativeEnvironment):
46 """NativeEnvironment for Airflow task templates."""
47
48
49class SandboxedEnvironment(_AirflowEnvironmentMixin, jinja2.sandbox.SandboxedEnvironment):
50 """SandboxedEnvironment for Airflow task templates."""
51
52
53def ds_filter(value: datetime.date | datetime.time | None) -> str | None:
54 """Date filter."""
55 if value is None:
56 return None
57 return value.strftime("%Y-%m-%d")
58
59
60def ds_nodash_filter(value: datetime.date | datetime.time | None) -> str | None:
61 """Date filter without dashes."""
62 if value is None:
63 return None
64 return value.strftime("%Y%m%d")
65
66
67def ts_filter(value: datetime.date | datetime.time | None) -> str | None:
68 """Timestamp filter."""
69 if value is None:
70 return None
71 return value.isoformat()
72
73
74def ts_nodash_filter(value: datetime.date | datetime.time | None) -> str | None:
75 """Timestamp filter without dashes."""
76 if value is None:
77 return None
78 return value.strftime("%Y%m%dT%H%M%S")
79
80
81def ts_nodash_with_tz_filter(value: datetime.date | datetime.time | None) -> str | None:
82 """Timestamp filter with timezone."""
83 if value is None:
84 return None
85 return value.isoformat().replace("-", "").replace(":", "")
86
87
88FILTERS = {
89 "ds": ds_filter,
90 "ds_nodash": ds_nodash_filter,
91 "ts": ts_filter,
92 "ts_nodash": ts_nodash_filter,
93 "ts_nodash_with_tz": ts_nodash_with_tz_filter,
94}