Skip to content

Commit

Permalink
feat: don't subst env vars for system
Browse files Browse the repository at this point in the history
Signed-off-by: xxchan <[email protected]>
  • Loading branch information
xxchan committed Jun 28, 2024
1 parent 2a451cc commit 9515496
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 23 deletions.
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

**Breaking changes**:
* runner: `RecordOutput` is now returned by `Runner::run` (or `Runner::run_async`). This allows users to access the output of each record, or check whether the record is skipped.
* substitution: add a special variable `__NOW__` which will be replaced with the current Unix timestamp in nanoseconds.
* runner(substitution): add a special variable `__NOW__` which will be replaced with the current Unix timestamp in nanoseconds.
* runner(substitution): for `system` commands, we do not substitute environment variables any more, because the shell can do that. It's necessary to escape like `\\` any more. `$__TEST_DIR__`, and are still supported.

## [0.20.6] - 2024-06-21

Expand Down
21 changes: 13 additions & 8 deletions sqllogictest/src/runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -576,7 +576,7 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
expected: _,
loc: _,
} => {
let sql = match self.may_substitute(sql) {
let sql = match self.may_substitute(sql, true) {
Ok(sql) => sql,
Err(error) => {
return RecordOutput::Statement {
Expand Down Expand Up @@ -627,7 +627,7 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
return RecordOutput::Nothing;
}

let mut command = match self.may_substitute(command) {
let mut command = match self.may_substitute(command, false) {
Ok(command) => command,
Err(error) => {
return RecordOutput::System {
Expand Down Expand Up @@ -723,7 +723,7 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
expected,
loc: _,
} => {
let sql = match self.may_substitute(sql) {
let sql = match self.may_substitute(sql, true) {
Ok(sql) => sql,
Err(error) => {
return RecordOutput::Query {
Expand Down Expand Up @@ -1196,12 +1196,17 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {

/// Substitute the input SQL or command with [`Substitution`], if enabled by `control
/// substitution`.
fn may_substitute(&self, input: String) -> Result<String, AnyError> {
#[derive(thiserror::Error, Debug)]
#[error("substitution failed: {0}")]
struct SubstError(subst::Error);
///
/// If `subst_env_vars`, we will use the `subst` crate to support extensive substitutions, incl. `$NAME`, `${NAME}`, `${NAME:default}`.
/// The cost is that we will have to use escape characters, e.g., `\$` & `\\`.
///
/// Otherwise, we just do simple string substitution for `__TEST_DIR__` and `__NOW__`.
/// This is useful for `system` commands: The shell can do the environment variables, and we can write strings like `\n` without escaping.
fn may_substitute(&self, input: String, subst_env_vars: bool) -> Result<String, AnyError> {
if let Some(substitution) = &self.substitution {
subst::substitute(&input, substitution).map_err(|e| Arc::new(SubstError(e)) as AnyError)
substitution
.substitute(&input, subst_env_vars)
.map_err(|e| Arc::new(e) as AnyError)
} else {
Ok(input)
}
Expand Down
47 changes: 33 additions & 14 deletions sqllogictest/src/substitution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,25 +11,44 @@ pub(crate) struct Substitution {
test_dir: Arc<OnceLock<TempDir>>,
}

#[derive(thiserror::Error, Debug)]
#[error("substitution failed: {0}")]
pub(crate) struct SubstError(subst::Error);

impl Substitution {
pub fn substitute(&self, input: &str, subst_env_vars: bool) -> Result<String, SubstError> {
if !subst_env_vars {
Ok(input
.replace("$__TEST_DIR__", &self.test_dir())
.replace("$__NOW__", &self.now()))
} else {
subst::substitute(&input, self).map_err(SubstError)
}
}

fn test_dir(&self) -> String {
let test_dir = self
.test_dir
.get_or_init(|| tempdir().expect("failed to create testdir"));
test_dir.path().to_string_lossy().into_owned()
}

fn now(&self) -> String {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("failed to get current time")
.as_nanos()
.to_string()
}
}

impl<'a> subst::VariableMap<'a> for Substitution {
type Value = String;

fn get(&'a self, key: &str) -> Option<Self::Value> {
match key {
"__TEST_DIR__" => {
let test_dir = self
.test_dir
.get_or_init(|| tempdir().expect("failed to create testdir"));
test_dir.path().to_string_lossy().into_owned().into()
}

"__NOW__" => std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("failed to get current time")
.as_nanos()
.to_string()
.into(),

"__TEST_DIR__" => self.test_dir().into(),
"__NOW__" => self.now().into(),
key => Env.get(key),
}
}
Expand Down

0 comments on commit 9515496

Please sign in to comment.