sql-data-guard 0.1.4__tar.gz → 0.1.6__tar.gz

{sql_data_guard-0.1.4/src/sql_data_guard.egg-info → sql_data_guard-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.4
+Version: 0.1.6
 Summary: Safety Layer for LLM Database Interactions
 Author-email: Imperva - Threat Reseach Infra <ww.dis.imperva.threat-research-infra@thalesgroup.com>
 License: MIT License

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6}/pyproject.toml

@@ -9,7 +9,7 @@ package-dir = {"" = "src"}
 extra-index-url = ["https://pypi.org/simple"]
 [project]
 name = "sql-data-guard"
-version = "0.1.4"
+version = "0.1.6"
 dependencies = [
     "sqlglot"
 ]

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6}/src/sql_data_guard/restriction_verification.py

@@ -130,7 +130,7 @@ def _verify_restriction(
     # Handle IN condition correctly
     if isinstance(exp, expr.In):
         expr_values = [str(val.this) for val in exp.expressions]
-        return any(v in values for v in expr_values)
+        return all(v in values for v in expr_values)
 
     # Handle EQ (=) condition
     if isinstance(exp, expr.EQ) and isinstance(exp.right, expr.Condition):

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6}/src/sql_data_guard/sql_data_guard.py

@@ -79,7 +79,7 @@ def _verify_where_clause(
 ):
     where_clause = select_statement.find(expr.Where)
     if where_clause:
-        for sub in where_clause.find_all(expr.Subquery):
+        for sub in where_clause.find_all(expr.Subquery, expr.Exists):
             _verify_query_statement(sub.this, context)
     _verify_static_expression(select_statement, context)
     verify_restrictions(select_statement, context, from_tables)
@@ -136,12 +136,10 @@ def _verify_query_statement(query_statement: expr.Query, context: VerificationCo
         _add_table_alias(cte, context)
         _verify_query_statement(cte.this, context)
     from_tables = _verify_from_tables(context, query_statement)
-    if not context.can_fix:
-        return query_statement
-    _verify_select_clause(context, query_statement, from_tables)
-    _verify_where_clause(context, query_statement, from_tables)
-    _verify_order_by_clause(context, query_statement)
-    return query_statement
+    if context.can_fix:
+        _verify_select_clause(context, query_statement, from_tables)
+        _verify_where_clause(context, query_statement, from_tables)
+        _verify_sub_queries(context, query_statement)
 
 
 def _verify_from_tables(context, query_statement):
@@ -156,11 +154,12 @@ def _verify_from_tables(context, query_statement):
     return from_tables
 
 
-def _verify_order_by_clause(context, query_statement):
-    for exp_type in [expr.Order, expr.Offset]:
+def _verify_sub_queries(context, query_statement):
+    for exp_type in [expr.Order, expr.Offset, expr.Limit, expr.Group, expr.Having]:
         for exp in find_direct(query_statement, exp_type):
-            for sub in exp.find_all(expr.Subquery):
-                _verify_query_statement(sub.this, context)
+            if exp:
+                for sub in exp.find_all(expr.Subquery):
+                    _verify_query_statement(sub.this, context)
 
 
 def _verify_select_clause(
@@ -289,6 +288,8 @@ def _get_from_clause_tables(
 def _add_table_alias(exp: expr.Expression, context: VerificationContext):
     for table_alias in find_direct(exp, expr.TableAlias):
         if isinstance(table_alias, expr.TableAlias):
-            context.dynamic_tables[table_alias.alias_or_name] = {
-                col.alias_or_name for col in table_alias.columns
-            }
+            if len(table_alias.columns) > 0:
+                column_names = {col.alias_or_name for col in table_alias.columns}
+            else:
+                column_names = {c for c in exp.this.named_selects}
+            context.dynamic_tables[table_alias.alias_or_name] = column_names

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6/src/sql_data_guard.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.4
+Version: 0.1.6
 Summary: Safety Layer for LLM Database Interactions
 Author-email: Imperva - Threat Reseach Infra <ww.dis.imperva.threat-research-infra@thalesgroup.com>
 License: MIT License

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6}/test/test_sql_guard_curr_unit.py

@@ -701,10 +701,12 @@ class TestSQLOrderRestrictions:
         verify_sql_test(
             sql,
             config,
-            errors=set(),  # No errors expected as the restriction matches the IN clause
-            fix=None,  # No fix should be needed
+            errors={
+                "Missing restriction for table: orders column: account_id value: [124]"
+            },
+            fix="SELECT product_name FROM orders WHERE (account_id IN (123, 124, 125)) AND account_id = 124",
             cnn=cnn,
-            data=[("Product A",), ("Product B",), ("Product C",)],
+            data=[("Product B",)],
         )
 
     def test_id_greater_than_122_should_return_error(self, config):

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6}/test/test_sql_guard_unit.py

@@ -6,8 +6,8 @@ from typing import Generator
 
 import pytest
 
+from conftest import verify_sql_test
 from sql_data_guard import verify_sql
-from conftest import verify_sql_test, verify_sql_test_data
 
 
 def _get_resource(file_name: str) -> str:
@@ -354,6 +354,32 @@ class TestTrino:
         )
 
 
+class TestTrinoWithRestrictions:
+    @pytest.fixture(scope="class")
+    def config(self) -> dict:
+        return {
+            "tables": [
+                {
+                    "table_name": "accounts",
+                    "columns": ["id", "day", "product_name"],
+                    "restrictions": [
+                        {"column": "id", "value": 123},
+                    ],
+                }
+            ]
+        }
+
+    @pytest.mark.skip(reason="Bug in sqlglot. Waiting for fix")
+    def test_date_add(self, config):
+        verify_sql_test(
+            "SELECT id FROM accounts WHERE DATE(day) >= DATE_ADD('day', -7, CURRENT_DATE)",
+            config,
+            dialect="trino",
+            errors={"Missing restriction for table: accounts column: id value: 123"},
+            fix="SELECT id FROM accounts WHERE DATE(day) >= DATE_ADD('day', -7, CURRENT_DATE) AND id = 123",
+        )
+
+
 class TestRestrictionsWithDifferentDataTypes:
     @pytest.fixture(scope="class")
     def config(self) -> dict:

{sql_data_guard-0.1.4 → sql_data_guard-0.1.6}/test/test_sql_guard_updates_unit.py

@@ -705,6 +705,10 @@ class TestMultipleRestriction:
         verify_sql_test(
             """SELECT id FROM products1 WHERE id IN (324, 233)""",
             config,
+            errors={
+                "Missing restriction for table: products1 column: id value: [324, 224]"
+            },
+            fix="SELECT id FROM products1 WHERE (id IN (324, 233)) AND id IN (324, 224)",
             cnn=cnn,
             data=[["324"]],
         )