improve api

Author: rustyconover

README.md

@@ -1,6 +1,34 @@
 # [Query.Farm](https://query.farm) SQL Manipulation
 
-A Python library for intelligent SQL predicate manipulation using [SQLGlot](https://sqlglot.com/sqlglot.html). This library provides tools to safely remove specific predicates from `SQL WHERE` clauses and filter SQL statements based on column availability.
+A Python library for intelligent SQL predicate manipulation using [SQLGlot](https://sqlglot.org)
+
+### Column Filtering with Complex Expressions
+
+```python
+import sqlglot
+from query_farm_sql_manipulation import transforms
+
+sql = '''
+SELECT * FROM users
+WHERE age > 18
+  AND (status = 'active' OR role = 'admin')
+  AND department IN ('engineering', 'sales')
+'''
+
+# Parse the statement first
+statement = sqlglot.parse_one(sql, dialect="duckdb")
+
+# Only keep predicates involving 'age' and 'role'
+allowed_columns = {'age', 'role'}
+
+result = transforms.filter_column_references(
+    statement=statement,
+    selector=lambda col: col.name in allowed_columns,
+)
+
+# Result: SELECT * FROM users WHERE age > 18 AND role = 'admin'
+print(result.sql())
+```
 
 ## Features
 
@@ -44,19 +72,22 @@ transforms.remove_expression_part(target_predicate)
 print(statement.sql())
 ```
 
-### Column-Based Filtering
+### Column-Name Based Filtering
 
 ```python
+import sqlglot
 from query_farm_sql_manipulation import transforms
 
-# Filter SQL to only include predicates with allowed columns
+# Parse SQL statement first
 sql = 'SELECT * FROM data WHERE color = "red" AND size > 10 AND type = "car"'
+statement = sqlglot.parse_one(sql, dialect="duckdb")
+
+# Filter to only include predicates with allowed columns
 allowed_columns = {"color", "type"}
 
-filtered = transforms.filter_column_references_statement(
-    sql=sql,
-    allowed_column_names=allowed_columns,
-    dialect="duckdb"
+filtered = transforms.filter_column_references(
+    statement=statement,
+    selector=lambda col: col.name in allowed_columns,
 )
 
 # Result: SELECT * FROM data WHERE color = "red" AND type = "car"
@@ -67,7 +98,7 @@ print(filtered.sql())
 
 ### `remove_expression_part(child: sqlglot.Expression) -> None`
 
-Removes the specified expression from its parent, respecting logical structure.
+Removes the specified SQLGlot expression from its parent, respecting logical structure.
 
 **Parameters:**
 - `child`: The SQLGlot expression to remove
@@ -82,21 +113,40 @@ Removes the specified expression from its parent, respecting logical structure.
 - `NOT` expressions: Removes the entire NOT expression
 - `CASE` statements: Removes conditional branches
 
-### `filter_column_references_statement(*, sql: str, allowed_column_names: Container[str], dialect: str = "duckdb") -> sqlglot.Expression`
+### `filter_column_references(*, statement: sqlglot.Expression, selector: Callable[[sqlglot.expressions.Column], bool]) -> sqlglot.Expression`
 
-Filters a SQL statement to remove predicates containing columns not in the allowed set.
+Filters a SQL statement to remove predicates containing columns that don't match the selector criteria.
 
 **Parameters:**
-- `sql`: The SQL statement to filter
-- `allowed_column_names`: Container of column names that should be preserved
-- `dialect`: SQL dialect for parsing (default: "duckdb")
+- `statement`: The SQLGlot expression to filter
+- `selector`: A callable that takes a Column and returns True if it should be preserved, False if it should be removed
 
 **Returns:**
-- Filtered SQLGlot expression with non-allowed columns removed
+- Filtered SQLGlot expression with non-matching columns removed
 
 **Raises:**
 - `ValueError`: If a column can't be cleanly removed due to interactions with allowed columns
 
+### `where_clause_contents(statement: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression | None`
+
+Extracts the contents of the WHERE clause from a SQLGlot expression.
+
+**Parameters:**
+- `statement`: The SQLGlot expression to extract from
+
+**Returns:**
+- The contents of the WHERE clause, or None if no WHERE clause exists
+
+### `filter_predicates_with_right_side_column_references(statement: sqlglot.expressions.Expression) -> sqlglot.Expression`
+
+Filters out predicates that have column references on the right side of comparisons.
+
+**Parameters:**
+- `statement`: The SQLGlot expression to filter
+
+**Returns:**
+- Filtered SQLGlot expression with right-side column reference predicates removed
+
 ## Examples
 
 ### Complex Logic Handling
@@ -140,9 +190,15 @@ result = transforms.filter_column_references_statement(
 The library will raise `ValueError` when predicates cannot be safely removed:
 
 ```python
+import sqlglot
+from query_farm_sql_manipulation import transforms
+
 # This will raise ValueError because x = 1 is part of a larger expression
 sql = "SELECT * FROM data WHERE result = (x = 1)"
+statement = sqlglot.parse_one(sql, dialect="duckdb")
+
 # Cannot remove x = 1 because it's used as a value, not a predicate
+# This would raise ValueError if attempted
 ```
 
 ## Supported SQL Constructs

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "query-farm-sql-manipulation"
-version = "0.1.4"
+version = "0.1.5"
 description = "A Python library for intelligent SQL predicate manipulation using SQLGlot. This library provides tools to safely remove specific predicates from `SQL WHERE` clauses and filter SQL statements based on column availability."
 authors = [
     { name = "Rusty Conover", email = "rusty@query.farm" }

src/query_farm_sql_manipulation/test_transforms.py

@@ -44,7 +44,7 @@ def _remove_target_predicate(sql: str, target: str, expected: str) -> None:
     matching_predicates = [p for p in predicates if p == target_predicate]
 
     # Assert that we found at least one match
-    assert matching_predicates, f"No matching predicate found for: {target}"
+    assert matching_predicates, f"No matching predicate found for: {target} {expression}"
 
     # Remove all matching predicates
     for matched_predicate in matching_predicates:
@@ -145,6 +145,31 @@ def test_remove_logical_predicates_errors(sql: str, target: str) -> None:
         _remove_target_predicate(sql, target, "")
 
 
+@pytest.mark.parametrize(
+    "sql, expected", [("""x = 1""", "x = 1"), ("foo = bar and z = 4", "foo = bar AND z = 4")]
+)
+def test_where_clause_extract(sql: str, expected: str) -> None:
+    statement = sqlglot.parse_one(f'SELECT * FROM "data" WHERE {sql}')
+    extracted_where = transforms.where_clause_contents(statement)
+    assert extracted_where is not None, "Expected a WHERE clause to be present"
+    assert extracted_where.sql("duckdb") == expected
+
+
+@pytest.mark.parametrize(
+    "sql, expected",
+    [
+        ("""v1 >= v1 + 5 and z = 5""", "z = 5"),
+        ("""((v1 >= v1 + 5) or t1 = 5) and z = 5""", "(t1 = 5) AND z = 5"),
+    ],
+)
+def test_filter_predicates_with_right_side_column_references(sql: str, expected: str) -> None:
+    statement = sqlglot.parse_one(f'SELECT * FROM "data" WHERE {sql}')
+    updated = transforms.filter_predicates_with_right_side_column_references(statement)
+    extracted_where = transforms.where_clause_contents(updated)
+    assert extracted_where is not None, "Expected a WHERE clause to be present"
+    assert extracted_where.sql("duckdb") == expected
+
+
 @pytest.mark.parametrize(
     "column_names, sql, expected_sql",
     [
@@ -356,8 +381,10 @@ def test_filter_column_references(
 ) -> None:
     dialect = "duckdb"
     full_sql = f'SELECT * FROM "data" AS "data" WHERE {sql}'
-    result = transforms.filter_column_references_statement(
-        sql=full_sql, allowed_column_names=column_names
+    statement = sqlglot.parse_one(full_sql, dialect="duckdb")
+    result = transforms.filter_column_references(
+        statement=statement,
+        selector=lambda col: col.name in column_names,
     )
 
     optimized = sqlglot.optimizer.optimize(result, rules=RULES_WITHOUT_NORMALIZE, dialect=dialect)

src/query_farm_sql_manipulation/transforms.py

@@ -1,7 +1,8 @@
-from collections.abc import Container
+from collections.abc import Callable
 
 import sqlglot
 import sqlglot.expressions
+import sqlglot.optimizer.simplify
 
 
 def remove_expression_part(child: sqlglot.Expression) -> None:
@@ -65,36 +66,70 @@ def remove_expression_part(child: sqlglot.Expression) -> None:
         raise ValueError(f"Cannot remove child from parent of type {type(parent)} {parent.sql()}")
 
 
-def filter_column_references_statement(
-    *, sql: str, allowed_column_names: Container[str], dialect: str = "duckdb"
+def where_clause_contents(
+    statement: sqlglot.expressions.Expression,
+) -> sqlglot.expressions.Expression | None:
+    """
+    Extract the contents of the WHERE clause from a SQLGlot expression.
+    Args:
+        statement: The SQLGlot expression to extract from
+    Returns:
+        The contents of the WHERE clause, or None if no WHERE clause exists
+    """
+    where_clause = statement.find(sqlglot.expressions.Where)
+    if where_clause is None:
+        return None
+    return where_clause.this
+
+
+def filter_predicates_with_right_side_column_references(
+    statement: sqlglot.expressions.Expression,
+) -> sqlglot.Expression:
+    # Need to simplify the statement to move the column references to the
+    # left side by default.
+    statement = sqlglot.optimizer.simplify.simplify(statement)
+
+    # If there's no WHERE clause, nothing to filter
+    where_clause = statement.find(sqlglot.expressions.Where)
+    if where_clause is None:
+        return statement
+    assert where_clause is not None
+
+    for predicate in where_clause.find_all(sqlglot.expressions.Predicate):
+        assert predicate is not None
+        if predicate.right.find(sqlglot.expressions.Column):
+            remove_expression_part(predicate)
+
+    return statement
+
+
+def filter_column_references(
+    *,
+    statement: sqlglot.expressions.Expression,
+    selector: Callable[[sqlglot.expressions.Column], bool],
 ) -> sqlglot.Expression:
     """
     Filter SQL statement to remove predicates with columns not in allowed_column_names.
 
     Args:
         sql: The SQL statement to filter
-        allowed_column_names: Container of column names that should be preserved
-        dialect: The SQL dialect to use for parsing (default is "duckdb")
+        selector: A callable that determines if a column should be preserved.
+                  It should return True for columns that are allowed, and False for those to be removed.
 
     Returns:
         Filtered SQLGlot expression with non-allowed columns removed
 
     Raises:
         ValueError: If a column can't be cleanly removed due to interactions with allowed columns
     """
-    # Parse and optimize the statement for predictable traversal
-    statement = sqlglot.parse_one(sql, dialect=dialect)
-
     # If there's no WHERE clause, nothing to filter
     where_clause = statement.find(sqlglot.expressions.Where)
     if where_clause is None:
         return statement
 
     # Find all column references not in allowed_column_names
     column_refs_to_remove = [
-        col
-        for col in where_clause.find_all(sqlglot.expressions.Column)
-        if col.name not in allowed_column_names
+        col for col in where_clause.find_all(sqlglot.expressions.Column) if not selector(col)
     ]
 
     # Process each column reference that needs to be removed
@@ -104,7 +139,7 @@ def filter_column_references_statement(
         closest_expression = _find_closest_removable_expression(column_ref)
 
         # Check if removing this expression would affect allowed columns
-        if _can_safely_remove_expression(closest_expression, allowed_column_names):
+        if _can_safely_remove_expression(closest_expression, selector):
             remove_expression_part(closest_expression)
         else:
             raise ValueError(
@@ -132,14 +167,16 @@ def _find_closest_removable_expression(
 
 
 def _can_safely_remove_expression(
-    expression: sqlglot.expressions.Expression, allowed_column_names: Container[str]
+    expression: sqlglot.expressions.Expression,
+    selector: Callable[[sqlglot.expressions.Column], bool],
 ) -> bool:
     """
     Check if an expression can be safely removed without affecting allowed columns.
 
     Args:
         expression: The expression to check
-        allowed_column_names: Container of allowed column names
+        selector: A callable that determines if a column should be preserved.
+                  It should return True for columns that are allowed, and False for those to be removed.
 
     Returns:
         True if the expression can be safely removed, False otherwise
@@ -161,9 +198,7 @@ def _can_safely_remove_expression(
 
     # Check if this expression references any allowed columns
     allowed_columns_referenced = [
-        col.name
-        for col in expression.find_all(sqlglot.expressions.Column)
-        if col.name in allowed_column_names
+        col.name for col in expression.find_all(sqlglot.expressions.Column) if selector(col)
     ]
 
     # If there are no allowed columns referenced, it's safe to remove