sql-data-guard 0.2.0__tar.gz → 0.2.1__tar.gz

{sql_data_guard-0.2.0/src/sql_data_guard.egg-info → sql_data_guard-0.2.1}/PKG-INFO

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

{sql_data_guard-0.2.0 → sql_data_guard-0.2.1}/pyproject.toml

@@ -9,7 +9,7 @@ package-dir = {"" = "src"}
 extra-index-url = ["https://pypi.org/simple"]
 [project]
 name = "sql-data-guard"
-version = "0.2.0"
+version = "0.2.1"
 dependencies = [
     # lower bound: oldest version the test-suite passes on; upper bound: next untested major
     "sqlglot>=25.20.1,<31"

sql_data_guard-0.2.1/src/sql_data_guard/__init__.py

@@ -0,0 +1,4 @@
+from .config_suggestion import merge_config, suggest_config
+from .sql_data_guard import validate_config, verify_sql
+
+__all__ = ["verify_sql", "validate_config", "suggest_config", "merge_config"]

{sql_data_guard-0.2.0 → sql_data_guard-0.2.1}/src/sql_data_guard/config_suggestion.py

@@ -2,8 +2,9 @@
 Suggesting and applying allow-list changes to a configuration.
 
 ``suggest_config`` answers "what would this query need to be allowed?" and ``merge_config``
-applies such an answer to a configuration. Both only deal with the table/column allow-list
-(``tables``); restrictions, masks, function lists and limits are never touched.
+applies such an answer to a configuration. Both only deal with the allow-lists (``tables`` and
+their columns, and ``databases``); restrictions, masks, function lists, limits and catalogs are
+never touched.
 """
 
 import copy
@@ -43,8 +44,19 @@ def _has_column(entry: dict, column: str) -> bool:
 
 def validate_delta(delta: dict) -> None:
     """Raises ValueError unless ``delta`` has the shape produced by ``suggest_config``."""
-    if not isinstance(delta, dict) or set(delta) != {"tables"} or not isinstance(delta["tables"], list):
-        raise ValueError("A config delta must be an object with a single 'tables' list.")
+    if (
+        not isinstance(delta, dict)
+        or "tables" not in delta
+        or not set(delta) <= {"tables", "databases"}
+        or not isinstance(delta["tables"], list)
+    ):
+        raise ValueError(
+            "A config delta must be an object with a 'tables' list and optionally a 'databases' list."
+        )
+    if "databases" in delta:
+        databases = delta["databases"]
+        if not isinstance(databases, list) or not all(isinstance(d, str) and d for d in databases):
+            raise ValueError("'databases' must be a list of database names.")
     for table in delta["tables"]:
         if not isinstance(table, dict) or not set(table) <= _DELTA_TABLE_KEYS:
             raise ValueError(f"A delta table may only have the keys {sorted(_DELTA_TABLE_KEYS)}.")
@@ -60,14 +72,21 @@ def validate_delta(delta: dict) -> None:
 
 def merge_config(config: dict, delta: dict) -> dict:
     """
-    Returns a copy of ``config`` with the tables and columns of ``delta`` added.
+    Returns a copy of ``config`` with the tables, columns and databases of ``delta`` added.
 
     A delta table with the same name (and database) as a configured one adds its missing columns
-    to it; any other delta table is appended. Nothing else in the configuration changes, and the
-    arguments are not modified. Raises ValueError if ``delta`` is malformed.
+    to it; any other delta table is appended. Delta databases missing from ``databases`` are
+    appended (the list is created if the configuration has none). Nothing else in the
+    configuration changes, and the arguments are not modified. Raises ValueError if ``delta`` is
+    malformed.
     """
     validate_delta(delta)
     merged = copy.deepcopy(config)
+    if delta.get("databases"):
+        databases = merged.setdefault("databases", [])
+        for database in delta["databases"]:
+            if database.lower() not in {d.lower() for d in databases}:
+                databases.append(database)
     tables = merged.setdefault("tables", [])
     for table in delta["tables"]:
         existing = next(
@@ -92,12 +111,15 @@ def suggest_config(sql: str, config: dict, dialect: Optional[str] = None) -> Opt
     """
     Suggests the smallest allow-list addition that lets ``sql`` run as written.
 
-    Returns ``{"tables": [{"table_name", "database_name"?, "columns": [...]}, ...]}``, to be
-    applied with ``merge_config``; ``{"tables": []}`` if the query needs nothing added. Returns
-    None when widening the allow-list cannot make the query acceptable:
+    Returns ``{"tables": [{"table_name", "database_name"?, "columns": [...]}, ...]}``, plus
+    ``"databases": [...]`` when the configuration has a ``databases`` allow-list that lacks a
+    database the query uses, to be applied with ``merge_config``; ``{"tables": []}`` if the
+    query needs nothing added. Returns None when widening the allow-list cannot make the query
+    acceptable:
 
     - the query is rejected for another reason (DML/DDL, several statements, a blocked function,
-      a column that cannot be attributed to one table, ...);
+      a column that cannot be attributed to one table, a table of another catalog, an
+      unqualified table in strict mode, ...);
     - it needs a table without naming any column (``SELECT *`` or ``COUNT(*)`` on a table the
       configuration does not know), because a configured table must list its columns.
 
@@ -117,6 +139,7 @@ def _suggest(sql: str, config: dict, dialect: Optional[str]) -> Optional[dict]:
     working = copy.deepcopy(config)
     tables = working["tables"]
     delta_tables: dict = {}  # identity -> delta table
+    delta_databases: List[str] = []
 
     def delta_for(entry: dict) -> dict:
         key = _identity(entry["table_name"], entry.get("database_name"))
@@ -132,6 +155,15 @@ def _suggest(sql: str, config: dict, dialect: Optional[str]) -> Optional[dict]:
         _verify_sql(sql, working, dialect, violations)
         progress = False
         for v in violations:
+            if v["type"] == "database":
+                known = {d.lower() for d in working["databases"]}
+                if v["database_name"].lower() not in known:
+                    working["databases"].append(v["database_name"])
+                    delta_databases.append(v["database_name"])
+                    progress = True
+                continue
+            if v["type"] not in ("table", "column"):
+                return None  # e.g. a catalog mismatch: not an allow-list problem
             if v["table_name"] is None:
                 return None  # a column that belongs to no single table: nothing to suggest
             entry = _find_entry(tables, v["table_name"], v["database_name"])
@@ -161,7 +193,9 @@ def _suggest(sql: str, config: dict, dialect: Optional[str]) -> Optional[dict]:
     result = [t for t in delta_tables.values() if t["columns"]]
     if len(result) != len(delta_tables):
         return None
-    delta = {"tables": result}
+    delta: dict = {"tables": result}
+    if delta_databases:
+        delta["databases"] = delta_databases
     merged = merge_config(config, delta)
     verdict = _verify_sql(sql, merged, dialect)
     if not verdict["allowed"] and verdict["fixed"] is None:

sql_data_guard-0.2.1/src/sql_data_guard/restriction_validation.py

@@ -0,0 +1,260 @@
+import math
+
+import sqlglot
+import sqlglot.expressions as expr
+
+from .column_masking import validate_column_masks
+
+
+class UnsupportedRestrictionError(Exception):
+    pass
+
+
+def _validate_max_length(config: dict):
+    if "max_length" not in config:
+        return
+    value = config["max_length"]
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(
+            f"Invalid 'max_length'. Expected a positive integer. Received: {value!r}"
+        )
+
+
+def _is_scalar(value) -> bool:
+    """A restriction value: a string, a bool or a finite number."""
+    if isinstance(value, float):
+        return math.isfinite(value)
+    return isinstance(value, (bool, int, str))
+
+
+def _is_number(value) -> bool:
+    return _is_scalar(value) and not isinstance(value, (bool, str))
+
+
+def _is_column_name(name) -> bool:
+    """A plain (optionally quoted) column name, as it is pasted into the condition that is added."""
+    if not isinstance(name, str) or not name.strip():
+        return False
+    try:
+        statements = sqlglot.parse(name)
+    except sqlglot.errors.SqlglotError:
+        return False
+    if len(statements) != 1:  # `id; DROP TABLE x`
+        return False
+    parsed = statements[0]
+    return isinstance(parsed, expr.Column) and not parsed.table
+
+
+def _validate_columns(table: dict):
+    columns = table.get("columns")
+    if not isinstance(columns, list) or not columns:
+        raise ValueError(
+            "Each table must have a 'columns' key with valid column definitions."
+        )
+    if not all(isinstance(c, str) and c.strip() for c in columns):
+        raise ValueError(
+            f"The 'columns' of table '{table['table_name']}' must be non-empty strings."
+        )
+
+
+def _validate_restriction_shape(restriction, table_name: str):
+    """Checks the keys every restriction needs, whatever its operation."""
+    if not isinstance(restriction, dict):
+        raise ValueError(f"Each restriction of table '{table_name}' must be an object.")
+    if not _is_column_name(restriction.get("column")):
+        raise ValueError(
+            f"Each restriction of table '{table_name}' must have a 'column' with a column name. "
+            f"Received: {restriction.get('column')!r}"
+        )
+    operation = restriction.get("operation")
+    if operation is not None and not isinstance(operation, str):
+        raise UnsupportedRestrictionError(
+            f"Invalid restriction: 'operation={operation!r}' is not supported."
+        )
+    if operation in (None, "="):
+        # `value`, or `values` (the column may then equal any of them)
+        if "value" in restriction:
+            ok = _is_scalar(restriction["value"])
+        else:
+            values = restriction.get("values")
+            ok = (
+                isinstance(values, list)
+                and len(values) > 0
+                and all(_is_scalar(v) for v in values)
+            )
+        if not ok:
+            raise ValueError(
+                f"Invalid restriction for column '{restriction['column']}' in table '{table_name}'. "
+                "Expected a 'value' (a number, string or boolean) or a non-empty list of 'values'."
+            )
+
+
+def _non_empty_string(value) -> bool:
+    return isinstance(value, str) and bool(value.strip())
+
+
+def _validate_database_options(config: dict):
+    """`databases`, `strict_database` and `catalog_name` (top level and per table)."""
+    if not isinstance(config.get("strict_database", False), bool):
+        raise ValueError(
+            "Invalid 'strict_database'. Expected true or false. "
+            f"Received: {config['strict_database']!r}"
+        )
+    if "catalog_name" in config and not _non_empty_string(config["catalog_name"]):
+        raise ValueError(
+            f"Invalid 'catalog_name'. Expected a non-empty string. Received: {config['catalog_name']!r}"
+        )
+    if "databases" in config:
+        databases = config["databases"]
+        if not isinstance(databases, list) or not all(
+            _non_empty_string(d) for d in databases
+        ):
+            raise ValueError(
+                f"Invalid 'databases'. Expected a list of database names. Received: {databases!r}"
+            )
+
+
+def _validate_table_database(config: dict, table: dict):
+    name = table["table_name"]
+    if "catalog_name" in table and not _non_empty_string(table["catalog_name"]):
+        raise ValueError(f"The 'catalog_name' of table '{name}' must be a non-empty string.")
+    database_name = table.get("database_name")
+    if config.get("strict_database") and not _non_empty_string(database_name):
+        raise ValueError(
+            f"Table '{name}' must have a 'database_name' when 'strict_database' is set."
+        )
+    if database_name and "databases" in config:
+        if database_name.lower() not in {d.lower() for d in config["databases"]}:
+            raise ValueError(
+                f"The database '{database_name}' of table '{name}' is not in 'databases'."
+            )
+
+
+def _validate_force_limit(config: dict):
+    if "force_limit" not in config:
+        return
+    limit = config["force_limit"]
+    if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0:
+        raise ValueError(
+            f"Invalid 'force_limit'. Expected a positive integer. Received: {limit!r}"
+        )
+
+
+def _validate_function_list(config: dict, key: str):
+    if key not in config:
+        return
+    names = config[key]
+    if not isinstance(names, list) or not all(
+        isinstance(n, str) and n.strip() for n in names
+    ):
+        raise ValueError(
+            f"Invalid '{key}'. Expected a list of function names. Received: {names!r}"
+        )
+
+
+def validate_restrictions(config: dict):
+    """
+    Validates the restrictions in the configuration to ensure only supported operations are used.
+
+    Args:
+        config (dict): The configuration containing the restrictions to validate.
+
+    Raises:
+        UnsupportedRestrictionError: If an unsupported restriction operation is found.
+        ValueError: If there are no tables in the configuration.
+    """
+    supported_operations = [
+        "=",
+        ">",
+        "<",
+        ">=",
+        "<=",
+        "BETWEEN",
+        "IN",
+    ]  # Allowed operations
+    _validate_max_length(config)
+    _validate_database_options(config)
+    _validate_force_limit(config)
+    _validate_function_list(config, "blocked_functions")
+    _validate_function_list(config, "allowed_functions")
+    # Ensure 'tables' exists in config and is not empty
+    tables = config.get("tables", [])
+    # Check if tables are empty
+    if not tables:
+        raise ValueError("Configuration must contain at least one table.")
+    if not isinstance(tables, list):
+        raise ValueError("'tables' must be a list.")
+
+    seen = set()
+    for table in tables:
+        if not isinstance(table, dict):
+            raise ValueError("Each table must be an object.")
+        # Ensure that 'table_name' exists in each table
+        if "table_name" not in table:
+            raise ValueError("Each table must have a 'table_name' key.")
+        if not isinstance(table["table_name"], str) or not table["table_name"].strip():
+            raise ValueError("Each 'table_name' must be a non-empty string.")
+        database_name = table.get("database_name")
+        if database_name is not None and not isinstance(database_name, str):
+            raise ValueError(
+                f"The 'database_name' of table '{table['table_name']}' must be a string."
+            )
+        _validate_table_database(config, table)
+        key = (
+            table["table_name"].lower(),
+            (database_name or "").lower(),
+            str(table.get("catalog_name") or "").lower(),
+        )
+        if key in seen:
+            raise ValueError(
+                f"Table '{table['table_name']}' is defined more than once."
+            )
+        seen.add(key)
+        _validate_columns(table)
+
+        restrictions = table.get("restrictions", [])
+        if not isinstance(restrictions, list):
+            raise ValueError(
+                f"The 'restrictions' of table '{table['table_name']}' must be a list."
+            )
+
+        for restriction in restrictions:
+            _validate_restriction_shape(restriction, table["table_name"])
+            operation = restriction.get("operation")
+            if operation == "BETWEEN":
+                values = restriction.get("values")
+                if not (
+                    isinstance(values, list)
+                    and len(values) == 2
+                    and all(_is_number(v) for v in values)
+                    and values[0] < values[1]
+                ):
+                    raise ValueError(
+                        f"Invalid 'BETWEEN' format. Expected list of two numeric values where min < max. Received: {values}"
+                    )
+
+            elif operation == "IN":
+                values = restriction.get("values")
+                if not (
+                    isinstance(values, list)
+                    and len(values) > 0
+                    and all(_is_scalar(v) and not isinstance(v, bool) for v in values)
+                ):
+                    raise ValueError(
+                        f"Invalid 'IN' format. Expected a non-empty list of numbers or strings. Received: {values}"
+                    )
+
+            elif operation in (">", "<", ">=", "<="):
+                # comparison operators are evaluated numerically
+                value = restriction.get("value")
+                if not _is_number(value):
+                    raise ValueError(
+                        f"Invalid restriction value type for column '{restriction['column']}' in table '{table['table_name']}'. Expected a numeric value."
+                    )
+
+            elif operation and operation.lower() not in supported_operations:
+                raise UnsupportedRestrictionError(
+                    f"Invalid restriction: 'operation={operation}' is not supported."
+                )
+
+    validate_column_masks(config)

{sql_data_guard-0.2.0 → sql_data_guard-0.2.1}/src/sql_data_guard/restriction_verification.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import sqlglot
 import sqlglot.expressions as expr
@@ -24,7 +24,13 @@ def verify_restrictions(
     else:
         and_exps = list(split_to_expressions(where_clause.this, expr.And))
     for c_t in context.config["tables"]:
-        for from_t in [t for t in from_tables if table_matches(t, c_t)]:
+        for from_t in [
+            t
+            for t in from_tables
+            # a CTE with the name of a configured table is not that table (its body is verified,
+            # with its restrictions, on its own)
+            if table_matches(t, c_t) and (t.db or t.name not in context.cte_names)
+        ]:
             for r in c_t.get("restrictions", []):
                 found = False
                 for sub_exp in and_exps:
@@ -34,7 +40,10 @@ def verify_restrictions(
                 if not found:
                     if from_t.alias:
                         t_prefix = f"{from_t.alias}."
-                    elif len([t for t in from_tables if same_name(t.name, from_t.name)]) > 1:
+                    elif (
+                        len([t for t in from_tables if same_name(t.name, from_t.name)])
+                        > 1
+                    ):
                         t_prefix = f"{from_t.name}."
                     else:
                         t_prefix = ""
@@ -87,11 +96,15 @@ def _create_new_condition(
         operand = f"({', '.join(_format_value(v, dialect) for v in values)})"
     else:
         operator = operation if operation in _COMPARISON_OPERATORS else "="
-        operand = (
-            _format_value(restriction["value"], dialect)
-            if "value" in restriction
-            else str(restriction["values"])[1:-1]
+        values = (
+            [restriction["value"]] if "value" in restriction else restriction["values"]
         )
+        if operator == "=" and len(values) > 1:
+            # `=` with several allowed values: the column may equal any of them
+            operator = "IN"
+            operand = f"({', '.join(_format_value(v, dialect) for v in values)})"
+        else:
+            operand = _format_value(values[0], dialect)
     new_condition = sqlglot.parse_one(
         f"{table_prefix}{restriction['column']} {operator} {operand}",
         dialect=dialect,
@@ -143,54 +156,112 @@ def _verify_restriction(
     ):
         return False
 
-    if exp.this.table and from_table.alias and not same_name(exp.this.table, from_table.alias):
+    if (
+        exp.this.table
+        and from_table.alias
+        and not same_name(exp.this.table, from_table.alias)
+    ):
         return False
-    if exp.this.table and not from_table.alias and not same_name(exp.this.table, from_table.name):
+    if (
+        exp.this.table
+        and not from_table.alias
+        and not same_name(exp.this.table, from_table.name)
+    ):
         return False
 
     values = _get_restriction_values(restriction)  # Get correct restriction values
+    operation = restriction.get("operation")
 
-    # Handle IN condition correctly
+    # Only literal values count: a column (`region = eu`, where `eu` is a column of a derived
+    # table), a sub-query (`IN (SELECT 999)`) or an expression could evaluate to anything.
     if isinstance(exp, expr.In):
-        expr_values = [str(val.this) for val in exp.expressions]
-        return all(v in values for v in expr_values)
+        if exp.args.get("query") or exp.args.get("unnest") or not exp.expressions:
+            return False
+        expr_values = [_literal_value(val) for val in exp.expressions]
+        return all(
+            v is not None and _value_allowed(operation, values, v) for v in expr_values
+        )
 
-    # Handle EQ (=) condition
-    if isinstance(exp, expr.EQ) and isinstance(exp.right, expr.Condition):
-        return str(exp.right.this) in values
+    if isinstance(exp, expr.EQ):
+        value = _literal_value(exp.right)
+        return value is not None and _value_allowed(operation, values, value)
 
     if isinstance(exp, expr.Between):
-        low = _to_number(exp.args["low"].this)
-        high = _to_number(exp.args["high"].this)
-        restriction_values = [_to_number(v) for v in values]
-        if len(values) == 2 and None not in [low, high, *restriction_values]:
-            return restriction_values[0] <= low and high <= restriction_values[1]
+        low = _to_number(_literal_value(exp.args["low"]))
+        high = _to_number(_literal_value(exp.args["high"]))
+        if low is None or high is None:
+            return False
+        if operation == "BETWEEN":
+            low_limit, high_limit = (_to_number(v) for v in values)
+            return (
+                low_limit is not None
+                and high_limit is not None
+                and low_limit <= low
+                and high <= high_limit
+            )
+        if operation in (">", ">="):
+            return _value_allowed(operation, values, str(low))
+        if operation in ("<", "<="):
+            return _value_allowed(operation, values, str(high))
         return False
 
-    if isinstance(exp, (expr.LT, expr.LTE, expr.GT, expr.GTE)) and isinstance(
-        exp.right, expr.Condition
-    ):
-        if restriction.get("operation") not in _COMPARISON_OPERATORS:
+    if isinstance(exp, (expr.LT, expr.LTE, expr.GT, expr.GTE)):
+        if operation not in _COMPARISON_OPERATORS:
             return False
-        if len(values) != 1:
+        upper = isinstance(exp, (expr.LT, expr.LTE))  # `col < x` / `col <= x`
+        if upper != (operation in ("<", "<=")):
             return False
-        actual = _to_number(exp.right.this)
+        bound = _to_number(_literal_value(exp.right))
         limit = _to_number(values[0])
-        if actual is None or limit is None:
-            return False
-        if isinstance(exp, expr.LT) and restriction["operation"] == "<":
-            return actual < limit
-        elif isinstance(exp, expr.LTE) and restriction["operation"] == "<=":
-            return actual <= limit
-        elif isinstance(exp, expr.GT) and restriction["operation"] == ">":
-            return actual > limit
-        elif isinstance(exp, expr.GTE) and restriction["operation"] == ">=":
-            return actual >= limit
-        else:
+        if bound is None or limit is None:
             return False
+        # the query's bound must be at least as strict as the restriction's
+        if operation == "<" and isinstance(exp, expr.LTE):
+            return bound < limit
+        if operation == ">" and isinstance(exp, expr.GTE):
+            return bound > limit
+        return bound <= limit if upper else bound >= limit
     return False
 
 
+def _literal_value(node: Optional[expr.Expression]) -> Optional[str]:
+    """
+    The value of a literal operand, as a string (`'eu'`, `123`, `-5`, `TRUE`), or None if the
+    operand is anything else (a column, a sub-query, a function or other expression).
+    """
+    while isinstance(node, expr.Paren):
+        node = node.this
+    if isinstance(node, expr.Literal):
+        return node.this
+    if isinstance(node, expr.Boolean):
+        return str(node.this)
+    if (
+        isinstance(node, expr.Neg)
+        and isinstance(node.this, expr.Literal)
+        and not node.this.is_string
+    ):
+        return f"-{node.this.this}"
+    return None
+
+
+def _value_allowed(operation: Optional[str], values: List[str], value: str) -> bool:
+    """True if a column equal to `value` satisfies the restriction."""
+    if operation == "BETWEEN":
+        number, low, high = (_to_number(v) for v in (value, *values))
+        return None not in (number, low, high) and low <= number <= high
+    if operation in _COMPARISON_OPERATORS:
+        number, limit = _to_number(value), _to_number(values[0])
+        if number is None or limit is None:
+            return False
+        return {
+            "<": number < limit,
+            "<=": number <= limit,
+            ">": number > limit,
+            ">=": number >= limit,
+        }[operation]
+    return value in values
+
+
 def _get_restriction_values(restriction: dict) -> List[str]:
     if "values" in restriction:
         values = [str(v) for v in restriction["values"]]