store: Use double quotation format for parsed sql tables matched with prelude CTE

Fixes SQL-266

Signed-off-by: Tümay Tuzcu <[email protected]>

Author: Tümay Tuzcu

store/postgres/src/sql/formatter.rs

@@ -1,19 +1,29 @@
 use sqlparser::ast::{ObjectName, Statement, TableFactor, VisitMut, VisitorMut};
 use std::ops::ControlFlow;
 
+use super::Schema;
+
 pub struct Formatter<'a> {
     prelude: &'a str,
+    schema: &'a Schema,
 }
 
 impl<'a> Formatter<'a> {
-    pub fn new(prelude: &'a str) -> Self {
-        Self { prelude }
+    pub fn new(prelude: &'a str, schema: &'a Schema) -> Self {
+        Self { prelude, schema }
     }
 
     fn prepend_prefix_to_object_name_mut(&self, name: &mut ObjectName) {
-        let table_identifier: &mut Vec<_> = &mut name.0;
+        let table_identifier = &mut name.0;
         // remove all but the last identifier
         table_identifier.drain(0..table_identifier.len() - 1);
+
+        // 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) {
+                table_name.quote_style = Some('"');
+            }
+        }
     }
 
     pub fn format(&mut self, statement: &mut Statement) -> String {
@@ -42,9 +52,11 @@ impl VisitorMut for Formatter<'_> {
 
 #[cfg(test)]
 mod test {
+    use std::collections::HashSet;
+
     use super::*;
     use crate::sql::constants::SQL_DIALECT;
-    const CTE_PREFIX: &str = "WITH swap AS (
+    const CTE_PREFIX: &str = "WITH \"swap\" AS (
             SELECT 
             id,
             amount_in,
@@ -57,7 +69,17 @@ mod test {
 
     #[test]
     fn format_sql() {
-        let mut formatter = Formatter::new(CTE_PREFIX);
+        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 mut formatter = Formatter::new(CTE_PREFIX, &schema);
 
         let sql = "SELECT token_in, SUM(amount_in) AS amount FROM unknown.swap GROUP BY token_in";
 
@@ -71,7 +93,8 @@ mod test {
             result,
             format!(
                 "{} SELECT to_jsonb(sub.*) AS data FROM ( {} ) AS sub",
-                CTE_PREFIX, "SELECT token_in, SUM(amount_in) AS amount FROM swap GROUP BY token_in"
+                CTE_PREFIX,
+                "SELECT token_in, SUM(amount_in) AS amount FROM \"swap\" GROUP BY token_in"
             )
         );
     }

store/postgres/src/sql/parser.rs

@@ -84,7 +84,7 @@ impl Parser {
         let mut validator = Validator::new(&self.schema);
         validator.validate_statements(&statements)?;
 
-        let mut formatter = Formatter::new(&self.prelude);
+        let mut formatter = Formatter::new(&self.prelude, &self.schema);
 
         let statement = statements
             .get_mut(0)
@@ -131,34 +131,34 @@ mod test {
 
     const SQL_QUERY: &str = "
         with tokens as (
-            select * from (values
+            select * from (values 
             ('0x0000000000000000000000000000000000000000','ETH','Ethereum',18),
             ('0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48','USDC','USD Coin',6)
             ) as t(address,symbol,name,decimals)
         )
 
-        select
+        select 
         date,
         t.symbol,
         SUM(amount)/pow(10,t.decimals) as amount
-        from (select
+        from (select 
         date(to_timestamp(block_timestamp) at time zone 'utc') as date,
         token,
         amount
-        from swap_multi as sm
+        from swap_multi as sm 
         ,unnest(sm.amounts_in,sm.tokens_in) as smi(amount,token)
         union all
-        select
+        select 
         date(to_timestamp(block_timestamp) at time zone 'utc') as date,
         token,
         amount
-        from sgd1.swap_multi as sm
+        from sgd1.swap_multi as sm 
         ,unnest(sm.amounts_out,sm.tokens_out) as smo(amount,token)
         ) as tp
         inner join tokens as t on t.address = '0x' || encode(tp.token,'hex')
         group by tp.date,t.symbol,t.decimals
-        order by tp.date desc ,amount desc
-
+        order by tp.date desc ,amount desc 
+        
         ";
 
     fn test_layout() -> Layout {
@@ -185,7 +185,7 @@ mod test {
         assert_eq!(
             query,
             r#"WITH "swap_multi" AS (SELECT concat('0x', encode("id", 'hex')) AS "id", concat('0x', encode("sender", 'hex')) AS "sender", "amounts_in", "tokens_in", "amounts_out", "tokens_out", "referral_code", "block_number", "block_timestamp", concat('0x', encode("transaction_hash", 'hex')) AS "transaction_hash", "block$" FROM "sgd0815"."swap_multi"),
-"token" AS (SELECT "id", concat('0x', encode("address", 'hex')) AS "address", "symbol", "name", "decimals", "block_range" FROM "sgd0815"."token" WHERE "block_range" @> 2147483647) SELECT to_jsonb(sub.*) AS data FROM ( WITH tokens AS (SELECT * FROM (VALUES ('0x0000000000000000000000000000000000000000', 'ETH', 'Ethereum', 18), ('0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', 'USDC', 'USD Coin', 6)) AS t (address, symbol, name, decimals)) SELECT date, t.symbol, SUM(amount) / pow(10, t.decimals) AS amount FROM (SELECT date(to_timestamp(block_timestamp) AT TIME ZONE 'utc') AS date, token, amount FROM swap_multi AS sm, UNNEST(sm.amounts_in, sm.tokens_in) AS smi (amount, token) UNION ALL SELECT date(to_timestamp(block_timestamp) AT TIME ZONE 'utc') AS date, token, amount FROM swap_multi AS sm, UNNEST(sm.amounts_out, sm.tokens_out) AS smo (amount, token)) AS tp JOIN tokens AS t ON t.address = '0x' || encode(tp.token, 'hex') GROUP BY tp.date, t.symbol, t.decimals ORDER BY tp.date DESC, amount DESC ) AS sub"#
+"token" AS (SELECT "id", concat('0x', encode("address", 'hex')) AS "address", "symbol", "name", "decimals", "block_range" FROM "sgd0815"."token" WHERE "block_range" @> 2147483647) SELECT to_jsonb(sub.*) AS data FROM ( WITH tokens AS (SELECT * FROM (VALUES ('0x0000000000000000000000000000000000000000', 'ETH', 'Ethereum', 18), ('0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', 'USDC', 'USD Coin', 6)) AS t (address, symbol, name, decimals)) SELECT date, t.symbol, SUM(amount) / pow(10, t.decimals) AS amount FROM (SELECT date(to_timestamp(block_timestamp) AT TIME ZONE 'utc') AS date, token, amount FROM "swap_multi" AS sm, UNNEST(sm.amounts_in, sm.tokens_in) AS smi (amount, token) UNION ALL SELECT date(to_timestamp(block_timestamp) AT TIME ZONE 'utc') AS date, token, amount FROM "swap_multi" AS sm, UNNEST(sm.amounts_out, sm.tokens_out) AS smo (amount, token)) AS tp JOIN tokens AS t ON t.address = '0x' || encode(tp.token, 'hex') GROUP BY tp.date, t.symbol, t.decimals ORDER BY tp.date DESC, amount DESC ) AS sub"#
         );
     }
 }