SQLDatabase schema prefixing changes FROM/JOIN text inside SQL string literals
What happens
When SQLDatabase is configured with a schema, run_sql() rewrites FROM and JOIN matches throughout the query string, including inside quoted SQL literals. This changes returned data, even though the actual table reference is prefixed correctly. SQL comments are changed too.
Reproduction
With an attached SQLite schema named tenant and a tenant.users table containing ('Ada'):
from sqlalchemy import create_engine, event, text
from llama_index.core.utilities.sql_wrapper import SQLDatabase
engine = create_engine("sqlite:///:memory:")
@event.listens_for(engine, "connect")
def attach_schema(connection, record):
connection.execute("ATTACH DATABASE ':memory:' AS tenant")
with engine.begin() as connection:
connection.execute(text("CREATE TABLE tenant.users (name TEXT)"))
connection.execute(text("INSERT INTO tenant.users VALUES ('Ada')"))
db = SQLDatabase(engine, schema="tenant")
print(db.run_sql("SELECT 'FROM users', name FROM users")[0])Actual: [('FROM tenant.users', 'Ada')]
Expected: [('FROM users', 'Ada')]
The same result reproduced twice on the current source. _add_schema_prefix() also rewrites text in -- JOIN ignored and /* FROM hidden */ comments.
Cause and possible fix
_add_schema_prefix() applies re.sub() over the full SQL string, rather than just executable SQL spans. Masking quoted strings, quoted identifiers and comments while finding table references, then inserting schema prefixes at the original offsets, preserves the literal data. This can have focused regression tests for escaped quotes and comments as well as run_sql().
Source: run-llama/llama_index