fix sub queries in where and multiple joins

Author: ori-n

src/sql_data_guard/sql_data_guard.py

@@ -135,6 +135,10 @@ def _verify_query_statement(query_statement: expr.Query, context: VerificationCo
     for cte in query_statement.ctes:
         _add_table_alias(cte, context)
         _verify_query_statement(cte.this, context)
+    where_clause = query_statement.find(expr.Where)
+    if where_clause:
+        for sub in where_clause.find_all(expr.Subquery):
+            _verify_query_statement(sub.this, context)
     from_tables = _get_from_clause_tables(query_statement, context)
     for t in from_tables:
         found = False
@@ -251,7 +255,7 @@ def _get_from_clause_tables(
     """
     result = []
     from_clause = select_clause.find(expr.From)
-    join_clauses = list(select_clause.find_all(expr.Join))
+    join_clauses = select_clause.args.get("joins", [])
     for clause in [from_clause] + join_clauses:
         if clause:
             for t in find_direct(clause, expr.Table):

test/resources/orders_test.jsonl

@@ -77,4 +77,6 @@
 {"name": "inner_select_clause_restricted_table2", "sql": "SELECT (SELECT id FROM users LIMIT 1) AS id FROM orders WHERE id = 123", "errors": ["Table users is not allowed"]}
 {"name": "inner_select_clause_restricted_table3", "sql": "SELECT (SELECT id FROM users LIMIT 1) FROM orders WHERE id = 123", "errors": ["Table users is not allowed"]}
 {"name": "inner_select_clause_restricted_col1", "sql": "SELECT (SELECT col1, id FROM orders WHERE id = 123) FROM orders WHERE id = 123", "errors": ["Column col1 is not allowed. Column removed from SELECT clause"], "fix": "SELECT (SELECT id FROM orders WHERE id = 123) FROM orders WHERE id = 123", "data": [[123]]}
-{"name": "multiple_joins1", "sql": "SELECT id FROM orders AS o1 JOIN orders AS o2 JOIN users WHERE o1.id=123 AND o2.id=123", "errors": ["Table users is not allowed"]}
+{"name": "multiple_joins1", "sql": "SELECT id FROM orders AS o1 JOIN orders AS o2 JOIN users WHERE o1.id=123 AND o2.id=123", "errors": ["Table users is not allowed"]}
+{"name": "sub_query_in_where", "sql": "SELECT id FROM orders WHERE id = 123 AND id IN (SELECT id FROM orders WHERE id = 123)", "data": [[123]]}
+{"name": "sub_query_in_where_missing_restriction", "sql": "SELECT id FROM orders WHERE id = 123 AND id IN (SELECT id FROM orders)", "errors": ["Missing restriction for table: orders column: id value: 123"], "fix": "SELECT id FROM orders WHERE id = 123 AND id IN (SELECT id FROM orders WHERE id = 123)", "data": [[123]]}

test/test_sql_guard_unit.py

@@ -252,6 +252,7 @@ def test_complex_join(self, config, cnn):
         COUNT(o.order_id) AS order_count
     FROM orders o    
     JOIN products p ON o.product_id = p.product_id
+    WHERE o.account_id = 123
     GROUP BY o.account_id, p.product_name
 ),
 RankedProducts AS (
@@ -273,11 +274,16 @@ def test_complex_join(self, config, cnn):
 JOIN RankedProducts rp ON oc.product_name = rp.product_name
 WHERE oc.account_id IN (
     -- Filter accounts with at least 2 orders
-    SELECT account_id FROM orders GROUP BY account_id HAVING COUNT(order_id) >= 2
+    SELECT account_id FROM orders
+    WHERE account_id = 123
+    GROUP BY account_id HAVING COUNT(order_id) >= 2
 )
 ORDER BY oc.account_id, rp.product_rank;"""
-        verify_sql_test_data(
-            sql, config, cnn, [(123, "Laptop", 1, 1), (123, "Smartphone", 1, 1)]
+        verify_sql_test(
+            sql,
+            config,
+            cnn=cnn,
+            data=[(123, "Laptop", 1, 1), (123, "Smartphone", 1, 1)],
         )