How do I use QueryShield with LangChain SQL agents?

LangChain's SQLDatabaseChain and create_sql_agent are popular text-to-SQL primitives but ship with weak guardrails — the top_k row limit and a string-match denylist (QUERY_CHECKER) are bypassable. QueryShield slots in as a custom SQLDatabase wrapper or as a tool-level callback.

Pattern A — wrap the SQLDatabase:

from langchain_community.utilities import SQLDatabase
from queryshield.langchain import QueryShieldSQLDatabase

db = QueryShieldSQLDatabase.from_uri( DB_URI, policy="analytics-agent", api_key=os.getenv("QUERYSHIELD_API_KEY"), ) agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools") ```

Every db.run(query) call is validated; rejections raise an exception the agent's reasoning loop can recover from ("try a different query with a tenant_id filter").

Pattern B — callback-based:

from queryshield.langchain import QueryShieldCallback
agent_executor.invoke(
    {"input": user_query},
    config={"callbacks": [QueryShieldCallback(policy="analytics-agent", context=ctx)]},
)

The callback intercepts tool calls, validates SQL, and either approves, rewrites (injects required predicates), or rejects with an LLM-recoverable error message. Evidence log captures the original NL prompt, the LangChain trace, and the validation decision — auditable end-to-end.