sql-rewriter 0.1.2

ANTLR4-based SQL rewriting tool for permission management in LLM-generated SQL

SQL Rewriter - ANTLR4-based SQL Rewriting Tool

A Python package for SQL rewriting based on ANTLR4 grammar parsing. Currently provides the add_where_condition function, designed for permission management in LLM-generated SQL. Uses syntax parsing instead of regex matching, making it easier to handle various SQL formats generated by large language models.

How It Works

Uses ANTLR4 to parse SQL statements into syntax trees, then traverses the tree using the Visitor pattern to locate target table queries and intelligently add or merge WHERE conditions. If the original SQL already has a WHERE clause, it wraps the existing condition in parentheses before adding the new permission condition with AND (to prevent LLM SQL injection).

Installation

pip install sql-rewriter

Or install from source (if you want to modify the code):

git clone https://github.com/wangyang377/sql-rewriter.git
cd sql-rewriter
./scripts/generate_parser.sh  # Requires ANTLR4, see Development section below
pip install -e .

Usage

Basic Usage

from sql_rewriter import add_where_condition

# Add condition when there's no WHERE clause
sql = "SELECT * FROM users;"
new_sql = add_where_condition(sql, "age > 18", "users")
# Result: SELECT * FROM users WHERE age > 18;

# Append condition when WHERE clause already exists (wraps existing condition in parentheses)
sql = "SELECT * FROM users WHERE age > 18;"
new_sql = add_where_condition(sql, "status = 'active'", "users")
# Result: SELECT * FROM users WHERE (age > 18) AND status = 'active';

# JOIN queries - add condition only for specific table
sql = "SELECT * FROM users JOIN orders ON users.id = orders.user_id;"
new_sql = add_where_condition(sql, "users.status = 'active'", "users")
# Result: SELECT * FROM users JOIN orders ON users.id = orders.user_id WHERE users.status = 'active';

# Nested queries work too - precisely locates target table
sql = "SELECT * FROM orders WHERE status = 'pending' AND EXISTS (SELECT 1 FROM users WHERE users.id = orders.user_id);"
new_sql = add_where_condition(sql, "users.status = 'active'", "users")
# Result: SELECT * FROM orders WHERE status = 'pending' AND EXISTS (SELECT 1 FROM users WHERE users.id = orders.user_id AND users.status = 'active');

API Reference

add_where_condition(sql_text, new_condition, table_name=None)

Parameters:

• sql_text: Original SQL statement
• new_condition: WHERE condition to add (without the WHERE keyword)
• table_name: Target table name. Condition is only added if the FROM clause contains this table. If None, no processing is performed

Returns:

• Modified SQL statement (string)

Raises:

• ValueError: If SQL parsing fails

Development

If you clone the project from Git, you need to generate ANTLR parser code first:

# Install ANTLR4 (macOS)
brew install antlr

# Linux (Ubuntu/Debian)
sudo apt-get install antlr4

# Then generate code
./scripts/generate_parser.sh

Run tests:

cd tests
python test_parser.py