store: Use Layout for schema information

Author: lutter

store/postgres/src/sql/formatter.rs

@@ -1,16 +1,16 @@
 use sqlparser::ast::{ObjectName, Statement, TableFactor, VisitMut, VisitorMut};
 use std::ops::ControlFlow;
 
-use super::Schema;
+use crate::relational::{Layout, SqlName};
 
 pub struct Formatter<'a> {
     prelude: &'a str,
-    schema: &'a Schema,
+    layout: &'a Layout,
 }
 
 impl<'a> Formatter<'a> {
-    pub fn new(prelude: &'a str, schema: &'a Schema) -> Self {
-        Self { prelude, schema }
+    pub fn new(prelude: &'a str, layout: &'a Layout) -> Self {
+        Self { prelude, layout }
     }
 
     fn prepend_prefix_to_object_name_mut(&self, name: &mut ObjectName) {
@@ -20,7 +20,8 @@ impl<'a> Formatter<'a> {
 
         // Ensure schema tables has quotation to match up with prelude generated cte.
         if let Some(table_name) = table_identifier.last_mut() {
-            if self.schema.contains_key(&table_name.value) {
+            let sql_name = SqlName::verbatim(table_name.to_string());
+            if self.layout.table(&sql_name).is_some() {
                 table_name.quote_style = Some('"');
             }
         }
@@ -52,10 +53,19 @@ impl VisitorMut for Formatter<'_> {
 
 #[cfg(test)]
 mod test {
-    use std::collections::HashSet;
-
     use super::*;
-    use crate::sql::constants::SQL_DIALECT;
+    use crate::sql::{constants::SQL_DIALECT, test::make_layout};
+
+    const GQL: &str = "
+    type Swap @entity {
+        id: ID!
+        amountIn: BigDecimal!
+        amountOut: BigDecimal!
+        tokenIn: Bytes!
+        tokenOut: Bytes!
+    }
+    ";
+
     const CTE_PREFIX: &str = "WITH \"swap\" AS (
             SELECT
             id,
@@ -69,17 +79,9 @@ mod test {
 
     #[test]
     fn format_sql() {
-        let mut schema = Schema::new();
-        schema.insert(
-            "swap".to_string(),
-            HashSet::from_iter(
-                ["id", "amount_in", "amount_out", "token_in", "token_out"]
-                    .into_iter()
-                    .map(|s| s.to_string()),
-            ),
-        );
+        let layout = make_layout(GQL);
 
-        let mut formatter = Formatter::new(CTE_PREFIX, &schema);
+        let mut formatter = Formatter::new(CTE_PREFIX, &layout);
 
         let sql = "SELECT token_in, SUM(amount_in) AS amount FROM unknown.swap GROUP BY token_in";
 

store/postgres/src/sql/mod.rs

@@ -3,8 +3,27 @@ mod formatter;
 mod parser;
 mod validation;
 
-use std::collections::{HashMap, HashSet};
+pub use parser::Parser;
 
-pub(self) type Schema = HashMap<String, HashSet<String>>; // HashMap<Table, HashSet<Column>>
+#[cfg(test)]
+mod test {
+    use std::{collections::BTreeSet, sync::Arc};
 
-pub use parser::Parser;
+    use graph::{prelude::DeploymentHash, schema::InputSchema};
+
+    use crate::{
+        catalog::Catalog,
+        primary::{make_dummy_site, Namespace},
+        relational::Layout,
+    };
+
+    pub(crate) fn make_layout(gql: &str) -> Layout {
+        let subgraph = DeploymentHash::new("Qmasubgraph").unwrap();
+        let schema = InputSchema::parse_latest(gql, subgraph.clone()).unwrap();
+        let namespace = Namespace::new("sgd0815".to_string()).unwrap();
+        let site = Arc::new(make_dummy_site(subgraph, namespace, "anet".to_string()));
+        let catalog = Catalog::for_tests(site.clone(), BTreeSet::new()).unwrap();
+        let layout = Layout::new(site, &schema, catalog).unwrap();
+        layout
+    }
+}

store/postgres/src/sql/parser.rs

@@ -52,39 +52,23 @@ pub fn generate_table_prelude_from_layout(layout: &Layout) -> String {
 }
 
 pub struct Parser {
-    schema: super::Schema,
+    layout: Arc<Layout>,
     prelude: String,
 }
 
 impl Parser {
     pub fn new(layout: Arc<Layout>) -> Self {
-        Self {
-            schema: layout
-                .tables
-                .iter()
-                .filter(|(entity, _)| !entity.is_poi())
-                .map(|(_, table)| {
-                    (
-                        table.name.to_string(),
-                        table
-                            .columns
-                            .iter()
-                            .map(|column| column.name.to_string())
-                            .collect(),
-                    )
-                })
-                .collect(),
-            prelude: generate_table_prelude_from_layout(&layout),
-        }
+        let prelude = generate_table_prelude_from_layout(&layout);
+        Self { layout, prelude }
     }
 
     pub fn parse_and_validate(&self, sql: &str) -> Result<String> {
         let mut statements = sqlparser::parser::Parser::parse_sql(&SQL_DIALECT, sql)?;
 
-        let mut validator = Validator::new(&self.schema);
+        let mut validator = Validator::new(&self.layout);
         validator.validate_statements(&statements)?;
 
-        let mut formatter = Formatter::new(&self.prelude, &self.schema);
+        let mut formatter = Formatter::new(&self.prelude, &self.layout);
 
         let statement = statements
             .get_mut(0)

store/postgres/src/sql/validation.rs

@@ -2,7 +2,9 @@ use sqlparser::ast::{Expr, ObjectName, Query, SetExpr, Statement, TableFactor, V
 use std::result::Result;
 use std::{collections::HashSet, ops::ControlFlow};
 
-use super::{constants::ALLOWED_FUNCTIONS, Schema};
+use crate::relational::Layout;
+
+use super::constants::ALLOWED_FUNCTIONS;
 
 #[derive(thiserror::Error, Debug, PartialEq)]
 pub enum Error {
@@ -17,14 +19,14 @@ pub enum Error {
 }
 
 pub struct Validator<'a> {
-    schema: &'a Schema,
+    layout: &'a Layout,
     ctes: HashSet<String>,
 }
 
 impl<'a> Validator<'a> {
-    pub fn new(schema: &'a Schema) -> Self {
+    pub fn new(layout: &'a Layout) -> Self {
         Self {
-            schema,
+            layout,
             ctes: Default::default(),
         }
     }
@@ -54,9 +56,9 @@ impl<'a> Validator<'a> {
 
     fn validate_table_name(&mut self, name: &ObjectName) -> ControlFlow<Error> {
         if let Some(table_name) = name.0.last() {
-            let table_name = table_name.to_string().to_lowercase();
-            if !self.schema.contains_key(&table_name) && !self.ctes.contains(&table_name) {
-                return ControlFlow::Break(Error::UnknownTable(table_name));
+            let name = &table_name.value;
+            if !self.layout.table(name).is_some() && !self.ctes.contains(name) {
+                return ControlFlow::Break(Error::UnknownTable(name.to_string()));
             }
         }
         ControlFlow::Continue(())
@@ -114,38 +116,35 @@ impl Visitor for Validator<'_> {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::sql::constants::SQL_DIALECT;
-    use std::collections::{HashMap, HashSet};
+    use crate::sql::{constants::SQL_DIALECT, test::make_layout};
 
     fn validate(sql: &str) -> Result<(), Error> {
         let statements = sqlparser::parser::Parser::parse_sql(&SQL_DIALECT, sql).unwrap();
 
-        let schema: Schema = HashMap::from([(
-            "swap".to_owned(),
-            HashSet::from([
-                "vid".to_owned(),
-                "block$".to_owned(),
-                "id".to_owned(),
-                "sender".to_owned(),
-                "input_amount".to_owned(),
-                "input_token".to_owned(),
-                "amount_out".to_owned(),
-                "output_token".to_owned(),
-                "slippage".to_owned(),
-                "referral_code".to_owned(),
-                "block_number".to_owned(),
-                "block_timestamp".to_owned(),
-                "transaction_hash".to_owned(),
-            ]),
-        )]);
-
-        let mut validator = Validator::new(&schema);
+        const GQL: &str = "
+            type Swap @entity {
+                id: ID!
+                sender: Bytes!
+                inputAmount: BigDecimal!
+                inputToken: Bytes!
+                amountOut: BigDecimal!
+                outputToken: Bytes!
+                slippage: BigDecimal!
+                referralCode: String
+                blockNumber: Int!
+                blockTimestamp: Timestamp!
+                transactionHash: Bytes!
+            }";
+
+        let layout = make_layout(GQL);
+
+        let mut validator = Validator::new(&layout);
 
         validator.validate_statements(&statements)
     }
 
     #[test]
-    fn test_function_blacklisted() {
+    fn test_function_disallowed() {
         let result = validate(
             "
             SELECT
@@ -161,7 +160,7 @@ mod test {
     }
 
     #[test]
-    fn test_table_function_blacklisted() {
+    fn test_table_function_disallowed() {
         let result = validate(
             "
         SELECT
@@ -181,7 +180,7 @@ mod test {
     }
 
     #[test]
-    fn test_function_blacklisted_without_paranthesis() {
+    fn test_function_disallowed_without_paranthesis() {
         let result = validate(
             "
             SELECT
@@ -195,7 +194,7 @@ mod test {
     }
 
     #[test]
-    fn test_function_whitelisted() {
+    fn test_function_allowed() {
         let result = validate(
             "
             SELECT