From c6b63c8d6a088bd9ae44a29049059261f742ba33 Mon Sep 17 00:00:00 2001
From: Henry Fontanier <henryfontanier@Henrys-MacBook-Pro.local>
Date: Wed, 22 Nov 2023 10:25:45 +0100
Subject: [PATCH] allow templating

---
 core/src/blocks/database.rs | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/core/src/blocks/database.rs b/core/src/blocks/database.rs
index 2171fb89a07fb..9ddfbdbb32ce5 100644
--- a/core/src/blocks/database.rs
+++ b/core/src/blocks/database.rs
@@ -1,4 +1,5 @@
 use crate::blocks::block::{parse_pair, Block, BlockResult, BlockType, Env};
+use crate::data_sources::data_source;
 use crate::Rule;
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
@@ -7,6 +8,7 @@ use pest::iterators::Pair;
 use serde_json::{json, Value};
 use tokio::sync::mpsc::UnboundedSender;
 
+use super::block::replace_variables_in_string;
 use super::helpers::get_data_source_project;
 
 #[derive(Clone)]
@@ -82,19 +84,23 @@ impl Block for Database {
         env: &Env,
         _event_sender: Option<UnboundedSender<Value>>,
     ) -> Result<BlockResult> {
-        let project =
-            get_data_source_project(&self.workspace_id, &self.data_source_id, env).await?;
+        let workspace_id = replace_variables_in_string(&self.workspace_id, "workspace_id", env)?;
+        let data_source_id =
+            replace_variables_in_string(&self.data_source_id, "data_source_id", env)?;
+        let database_id = replace_variables_in_string(&self.database_id, "database_id", env)?;
+
+        let project = get_data_source_project(&workspace_id, &data_source_id, env).await?;
 
         let database = match env
             .store
-            .load_database(&project, &self.data_source_id, &self.database_id)
+            .load_database(&project, &data_source_id, &database_id)
             .await?
         {
             Some(d) => d,
             None => Err(anyhow!(
                 "Database `{}` not found in data source `{}`",
-                self.database_id,
-                self.data_source_id
+                database_id,
+                data_source_id
             ))?,
         };
 
@@ -105,8 +111,8 @@ impl Block for Database {
             Ok(r) => r,
             Err(e) => Err(anyhow!(
                 "Error querying database `{}` in data source `{}`: {}",
-                self.database_id,
-                self.data_source_id,
+                database_id,
+                data_source_id,
                 e
             ))?,
         };