Merge pull request #5 from davidsmfreire/refactor/fix-issues

Refactor | Fix issues

Author: davidsmfreire

sqlshield/src/finder/mod.rs

@@ -34,9 +34,12 @@ pub fn find_queries_in_code(
         match file_extension.as_ref() {
             "py" => (
                 tree_sitter_python::language(),
-                python::extract_query_from_node,
+                python::extract_query_string_from_node,
+            ),
+            "rs" => (
+                tree_sitter_rust::language(),
+                rust::extract_query_string_from_node,
             ),
-            "rs" => (tree_sitter_rust::language(), rust::extract_query_from_node),
             _ => panic!("{}", format!("File not supported {file_extension}")),
         };
 
@@ -78,39 +81,33 @@ fn find_queries_in_ast(
     let mut cursor = node.walk();
 
     for child in node.children(&mut cursor) {
-        let mut child_cursor = child.walk();
-
-        for component in child.children(&mut child_cursor) {
-            match query_extractor(&component, code) {
-                Some(string_content) => {
-                    // ! Duct tape
-                    if string_content.find("REPLACE").is_some() {
-                        continue;
-                    }
-                    // !
+        match query_extractor(&child, code) {
+            Some(string_content) => {
+                // ! Duct tape
+                if string_content.find("REPLACE").is_some() {
+                    find_queries_in_ast(&child, code, query_extractor, dialect, queries, None)
+                }
+                // !
 
-                    let query_at = component.start_position();
+                let query_at = child.start_position();
 
-                    let statements = sqlparser::parser::Parser::parse_sql(dialect, &string_content);
+                let statements = sqlparser::parser::Parser::parse_sql(dialect, &string_content);
 
-                    match statements {
-                        Ok(statements) => {
-                            queries.push(super::QueryInCode {
-                                line: query_at.row + 1,
-                                statements,
-                            });
-                        }
-                        Err(err) => {
-                            if verbose.unwrap_or(0) > 0 {
-                                eprintln!("{err} {string_content}");
-                            }
+                match statements {
+                    Ok(statements) => {
+                        queries.push(super::QueryInCode {
+                            line: query_at.row + 1,
+                            statements,
+                        });
+                    }
+                    Err(err) => {
+                        if verbose.unwrap_or(0) > 0 {
+                            eprintln!("{err} {string_content}");
                         }
                     }
                 }
-                None => continue,
             }
+            None => find_queries_in_ast(&child, code, query_extractor, dialect, queries, None),
         }
-
-        find_queries_in_ast(&child, code, query_extractor, dialect, queries, None);
     }
 }

sqlshield/src/finder/python.rs

@@ -1,6 +1,6 @@
 use regex::Regex;
 
-pub fn extract_query_from_node(node: &tree_sitter::Node, code: &[u8]) -> Option<String> {
+pub fn extract_query_string_from_node(node: &tree_sitter::Node, code: &[u8]) -> Option<String> {
     if node.kind() != "string" {
         return None;
     }

sqlshield/src/finder/rust.rs

@@ -1,6 +1,6 @@
 use regex::Regex;
 
-pub fn extract_query_from_node(node: &tree_sitter::Node, code: &[u8]) -> Option<String> {
+pub fn extract_query_string_from_node(node: &tree_sitter::Node, code: &[u8]) -> Option<String> {
     if node.kind() != "string_literal" {
         return None;
     }

sqlshield/src/schema/mod.rs

@@ -13,30 +13,17 @@ pub fn load_schema_from_file(file_path: &Path) -> Result<TablesAndColumns, Strin
     let file_extension = &file_path.extension().unwrap().to_string_lossy();
 
     match schema {
-        Ok(schema) => match file_extension.as_ref() {
-            "sql" => sql::load_schema(&schema),
-            _ => Err(format!("File not supported {file_extension}")),
-        },
+        Ok(schema) => load_schema(&schema, file_extension),
         Err(err) => Err(format!(
             "Could not open {:?} due to error: {err}",
             file_path
         )),
     }
 }
 
-trait LoadSchema {
-    fn to_schema(&self, schema_type: &str) -> Result<TablesAndColumns, String>;
-}
-
-impl LoadSchema for String {
-    fn to_schema(&self, schema_type: &str) -> Result<TablesAndColumns, String> {
-        load_schema(&self.as_bytes(), schema_type)
-    }
-}
-
 pub fn load_schema(schema: &[u8], schema_type: &str) -> Result<TablesAndColumns, String> {
     match schema_type.as_ref() {
-        "sql" => sql::load_schema(&schema),
-        _ => Err(format!("Schema type not supported {schema_type}")),
+        "sql" => sql::load_schema(schema),
+        _ => panic!("{}", format!("Schema type not supported {schema_type}")),
     }
 }

sqlshield/src/schema/sql.rs

@@ -2,8 +2,6 @@ use sqlparser::{ast::Statement, dialect::GenericDialect, parser::Parser};
 use std::collections::{HashMap, HashSet};
 
 pub fn load_schema(schema: &[u8]) -> Result<super::TablesAndColumns, String> {
-    // TODO return err instead of panics
-
     let schema_str = String::from_utf8_lossy(schema);
 
     let dialect = GenericDialect {};
@@ -25,3 +23,38 @@ pub fn load_schema(schema: &[u8]) -> Result<super::TablesAndColumns, String> {
     }
     return Ok(tables);
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::{HashMap, HashSet};
+
+    use super::load_schema;
+
+    #[test]
+    fn test_load_schema() {
+        let schema = "
+            CREATE TABLE users (
+                id INT PRIMARY KEY AUTO_INCREMENT,
+                name VARCHAR(255) NOT NULL
+            );
+            CREATE TABLE receipt (
+                id INT PRIMARY KEY AUTO_INCREMENT,
+                content VARCHAR(128),
+                user_id INT,
+                FOREIGN KEY (user_id) REFERENCES users(id)
+            );
+        ";
+
+        let expected_result: HashMap<String, HashSet<String>> = HashMap::from([
+            ("users".into(), HashSet::from(["id".into(), "name".into()])),
+            (
+                "receipt".into(),
+                HashSet::from(["id".into(), "content".into(), "user_id".into()]),
+            ),
+        ]);
+
+        let result = load_schema(&schema.as_bytes()).unwrap();
+
+        assert_eq!(result, expected_result);
+    }
+}