athena/
tera.rs

1//! Tera template engine setup and custom functions
2//!
3//! This module configures the Tera template engine for rendering SQL templates.
4//! It provides:
5//! - Template loading from the working directory
6//! - Custom Tera functions (e.g., `date_range`)
7//!
8//! # Custom Functions
9//!
10//! ## `date_range`
11//!
12//! Generates a list of dates between start and end dates (exclusive of end).
13//!
14//! ```sql
15//! {% set start_date = "2022-01-01" %}
16//! {% set end_date = "2022-01-05" %}
17//! {% for date in date_range(start=start_date, end=end_date) %}
18//!   PARTITION (date='{{ date }}')
19//! {% endfor %}
20//! ```
21//!
22//! This generates:
23//! ```text
24//! PARTITION (date='2022-01-01')
25//! PARTITION (date='2022-01-02')
26//! PARTITION (date='2022-01-03')
27//! PARTITION (date='2022-01-04')
28//! ```
29
30use chrono::{Duration, NaiveDate};
31use log::debug;
32use std::collections::HashMap;
33use std::ffi::OsStr;
34use std::path::PathBuf;
35use tera::{from_value, to_value, Error, Tera, Value};
36use walkdir::WalkDir;
37
38use crate::utils::is_dir;
39
40// Constants
41const DATE_FORMAT: &str = "%Y-%m-%d";
42const SQL_FILE_EXTENSION: &str = "sql";
43
44/// Get Tera template, load the template from working dir
45pub fn get_tera(target_path: PathBuf, working_dir: PathBuf) -> anyhow::Result<Tera> {
46    let is_dir = is_dir(&target_path);
47    let working_dir_str = working_dir
48        .to_str()
49        .ok_or_else(|| anyhow::anyhow!("working directory path is not valid UTF-8"))?;
50    let prefix = format!("{}/", working_dir_str);
51
52    let mut tera = Tera::default();
53
54    // Scan working_dir and adding .sql file as template
55    let templates = WalkDir::new(&working_dir)
56        .into_iter()
57        .filter_map(|e| e.ok())
58        .filter_map(|e| {
59            if e.path().extension() == Some(OsStr::new(SQL_FILE_EXTENSION)) {
60                Some(e)
61            } else {
62                None
63            }
64        })
65        .map(|e| {
66            let template_path = e.path().display().to_string();
67            let template_name = template_path.trim_start_matches(&prefix).to_string();
68            (template_path, Some(template_name))
69        })
70        .collect::<Vec<_>>();
71
72    debug!("Loaded: {:?}", templates);
73    tera.add_template_files(templates)?;
74
75    if !is_dir {
76        let template_path = target_path.display().to_string();
77        let template_name = target_path
78            .file_name()
79            .and_then(|n| n.to_str())
80            .ok_or_else(|| anyhow::anyhow!("could not get file name from target path"))?;
81        tera.add_template_file(template_path, Some(template_name))?;
82    }
83
84    // Register functions
85    tera.register_function("date_range", date_range);
86
87    Ok(tera)
88}
89
90fn get_native_date(key: &str, input: Option<&Value>) -> tera::Result<NaiveDate> {
91    if input.is_none() {
92        return Err(Error::msg(format!(
93            "Function `date_range` was called without a `{key}` argument",
94        )));
95    }
96
97    // Safety: We just checked that input is Some above
98    let input = input.unwrap();
99
100    match from_value::<String>(input.clone()) {
101        Ok(val) => match NaiveDate::parse_from_str(&val, DATE_FORMAT) {
102            Ok(v) => Ok(v),
103            Err(_) => {
104                Err(Error::msg(format!(
105                    "Function `date_range` received {key}={input} but `{key}` is invalid format ({DATE_FORMAT})",
106                )))
107            }
108        },
109        Err(_) => Err(Error::msg(format!("Function `date_range` received {key}={input} but it is not a string"))),
110    }
111}
112
113pub fn date_range(args: &HashMap<String, Value>) -> tera::Result<Value> {
114    let start = get_native_date("start", args.get("start"))?;
115    let end = get_native_date("end", args.get("end"))?;
116
117    let mut items = vec![];
118    let mut cursor = start;
119    while cursor < end {
120        items.push(cursor.format(DATE_FORMAT).to_string());
121        cursor += Duration::days(1);
122    }
123
124    to_value(items).map_err(|e| Error::msg(format!("failed to convert date range to value: {}", e)))
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn test_date_range() {
133        let mut args = HashMap::new();
134        args.insert("start".to_string(), to_value("2022-01-29").unwrap());
135        args.insert("end".to_string(), to_value("2022-02-02").unwrap());
136
137        let res = date_range(&args).unwrap();
138        assert_eq!(
139            res,
140            to_value(vec!["2022-01-29", "2022-01-30", "2022-01-31", "2022-02-01"]).unwrap()
141        );
142    }
143}
athena/tera.rs

athena/
tera.rs
