Move variables_table to Env

crates/query-engine/translation/src/translation/helpers.rs

@@ -14,14 +14,14 @@ pub struct Env<'a> {
     metadata: &'a metadata::Metadata,
     relationships: BTreeMap<String, models::Relationship>,
     mutations_version: &'a Option<metadata::mutations::MutationsVersion>,
+    variables_table: Option<sql::ast::TableReference>,
 }
 
 #[derive(Debug)]
 /// Stateful information changed throughout the translation process.
 pub struct State {
     native_queries: NativeQueries,
     global_table_index: TableAliasIndex,
-    variables_table: Option<sql::ast::TableReference>,
 }
 
 #[derive(Debug)]
@@ -94,11 +94,13 @@ impl<'a> Env<'a> {
         metadata: &'a metadata::Metadata,
         relationships: BTreeMap<String, models::Relationship>,
         mutations_version: &'a Option<metadata::mutations::MutationsVersion>,
+        variables_table: Option<sql::ast::TableReference>,
     ) -> Env<'a> {
         Env {
             metadata,
             relationships,
             mutations_version,
+            variables_table,
         }
     }
     /// Lookup a collection's information in the metadata.
@@ -182,6 +184,15 @@ impl<'a> Env<'a> {
                 type_name: scalar_type.clone(),
             })
     }
+
+    /// Try to get the variables table reference. This will fail if no variables were passed
+    /// as part of the query request.
+    pub fn get_variables_table(&self) -> Result<sql::ast::TableReference, Error> {
+        match &self.variables_table {
+            None => Err(Error::UnexpectedVariable),
+            Some(t) => Ok(t.clone()),
+        }
+    }
 }
 
 impl CollectionInfo {
@@ -219,7 +230,6 @@ impl Default for State {
         State {
             native_queries: NativeQueries::new(),
             global_table_index: TableAliasIndex(0),
-            variables_table: None,
         }
     }
 }
@@ -242,7 +252,6 @@ impl State {
                 let variables_table_alias = self.make_table_alias("%variables_table".to_string());
                 let table_reference =
                     sql::ast::TableReference::AliasedTable(variables_table_alias.clone());
-                self.variables_table = Some(table_reference.clone());
                 Some((
                     sql::helpers::from_variables(variables_table_alias),
                     table_reference,
@@ -251,15 +260,6 @@ impl State {
         }
     }
 
-    /// Try to get the variables table reference. This will fail if no variables were passed
-    /// as part of the query request.
-    pub fn get_variables_table(&self) -> Result<sql::ast::TableReference, Error> {
-        match &self.variables_table {
-            None => Err(Error::UnexpectedVariable),
-            Some(t) => Ok(t.clone()),
-        }
-    }
-
     /// Introduce a new native query to the generated sql.
     pub fn insert_native_query(
         &mut self,

crates/query-engine/translation/src/translation/mutation/translate.rs

@@ -19,7 +19,7 @@ pub fn translate(
     collection_relationships: BTreeMap<String, models::Relationship>,
     mutations_version: &Option<metadata::mutations::MutationsVersion>,
 ) -> Result<sql::execution_plan::Mutation, Error> {
-    let env = Env::new(metadata, collection_relationships, mutations_version);
+    let env = Env::new(metadata, collection_relationships, mutations_version, None);
 
     match operation {
         models::MutationOperation::Procedure {
@@ -237,7 +237,7 @@ fn translate_native_query(
 
     // add the procedure native query definition is a with clause.
     select.with = sql::ast::With {
-        common_table_expressions: crate::translation::query::native_queries::translate(state)?,
+        common_table_expressions: crate::translation::query::native_queries::translate(env, state)?,
     };
 
     // normalize ast

crates/query-engine/translation/src/translation/query/filtering.rs

@@ -380,7 +380,7 @@ fn translate_comparison_value(
         )),
         models::ComparisonValue::Variable { name: var } => Ok((
             values::translate_variable(
-                state.get_variables_table()?,
+                env.get_variables_table()?,
                 var.clone(),
                 &database::Type::ScalarType(typ.clone()),
             ),

crates/query-engine/translation/src/translation/query/mod.rs

@@ -21,9 +21,15 @@ pub fn translate(
     isolation_level: sql::ast::transaction::IsolationLevel,
     query_request: models::QueryRequest,
 ) -> Result<sql::execution_plan::ExecutionPlan<sql::execution_plan::Query>, Error> {
-    let env = Env::new(metadata, query_request.collection_relationships, &None);
     let mut state = State::new();
     let variables_from = state.make_variables_table(&query_request.variables);
+    let variables_table_ref = variables_from.clone().map(|(_, table_ref)| table_ref);
+    let env = Env::new(
+        metadata,
+        query_request.collection_relationships,
+        &None,
+        variables_table_ref,
+    );
     let (current_table, from_clause) = root::make_from_clause_and_reference(
         &query_request.collection,
         &query_request.arguments,
@@ -59,7 +65,7 @@ pub fn translate(
         state.make_table_alias("universe_agg".to_string()),
         // native queries if there are any
         sql::ast::With {
-            common_table_expressions: native_queries::translate(state)?,
+            common_table_expressions: native_queries::translate(&env, state)?,
         },
         select_set,
     );

crates/query-engine/translation/src/translation/query/native_queries.rs

@@ -4,14 +4,14 @@ use ndc_sdk::models;
 
 use super::values;
 use crate::translation::error::Error;
-use crate::translation::helpers::State;
+use crate::translation::helpers::{Env, State};
 use query_engine_metadata::metadata;
 use query_engine_sql::sql;
 
 /// Translate native queries collected in State by the translation proccess into CTEs.
-pub fn translate(state: State) -> Result<Vec<sql::ast::CommonTableExpression>, Error> {
+pub fn translate(env: &Env, state: State) -> Result<Vec<sql::ast::CommonTableExpression>, Error> {
     let mut ctes = vec![];
-    let variables_table = state.get_variables_table();
+    let variables_table = env.get_variables_table();
     let native_queries = state.get_native_queries();
 
     // for each found table expression