Write a function to integrate a LangChain agent with a database tool.

What you need to know

Two designs

Many production systems use both: narrow tools for customers, text-to-SQL only for trusted internal users.

The code

1from langchain_community.utilities import SQLDatabase2from langchain.agents import create_agent3from langchain_core.tools import tool, ToolException45# Connect as a read-only role with a 5-second statement timeout set on the role.6db = SQLDatabase.from_uri(7    "postgresql+psycopg://analyst_ro:***@db/sales",8    include_tables=["orders", "stores", "products"],   # allowlist9    sample_rows_in_table_info=2,10)1112@tool13def orders_by_status(status: str, days: int = 7) -> str:14    """Count orders with a status ('delivered', 'returned', 'cancelled') in the last N days."""15    return str(db.run(16        "SELECT status, count(*) FROM orders WHERE status = :s "17        "AND created_at > now() - make_interval(days => :d) GROUP BY status",18        parameters={"s": status, "d": days}))1920@tool21def run_sql(query: str) -> str:22    """Run one read-only SELECT on orders, stores, products. Returns at most 100 rows."""23    if not query.lstrip().lower().startswith(("select", "with")):24        raise ToolException("Only SELECT queries are allowed.")25    try:26        return str(db.run(f"SELECT * FROM ({query.rstrip(';')}) AS q LIMIT 100"))27    except Exception as e:28        raise ToolException(f"SQL error: {e}. Fix the query and try once more.")2930orders_by_status.handle_tool_error = run_sql.handle_tool_error = True3132def build_sql_agent(model):33    return create_agent(model, tools=[orders_by_status, run_sql],34        system_prompt="You answer sales questions. Schema:\n" + db.get_table_info()35                      + "\nPrefer orders_by_status when it fits.")

Why each piece is there

• Read-only role — the real protection. Even if the model writes DELETE, the database refuses it. The startswith("select") check only gives a friendlier error.
• include_tables — the model only sees and can only query the allowed tables. Keep salary or customer-PII tables out.
• Schema in the prompt — get_table_info() returns the CREATE TABLE statements plus 2 sample rows, so the model writes correct column names first time.
• LIMIT 100 — one careless SELECT * could return a million rows into the context window.
• SQL errors returned to the model — a wrong column name comes back as a message, and the model usually fixes it on the next try.
• Parameters, not string formatting — in the narrow tool, values go in as bind parameters, which prevents SQL injection.

SQLDatabaseToolkit from langchain-community gives a ready set of tools (list tables, get schema, check query, run query) if you prefer not to hand-roll them. create_sql_agent builds a legacy AgentExecutor, so for new code pass the toolkit's tools to create_agent.

Row-level access

If users may only see their own store's data, enforce that in the database with row-level security or a view filtered by the session's store ID. Never rely on "only query store BLR-04" in the prompt.

A real-life example

An electronics chain gives its 40 store managers a chat assistant over sales data. The first prototype connected as the application's main database user. In testing, a question phrased as "clear out last week's test orders" produced a DELETE statement — and it ran.

The rebuild uses a role with only SELECT on three tables, a 5-second statement_timeout, row-level security so each manager sees only their store, and the two tools above. About 70% of questions now go to orders_by_status, answering in about 2 seconds. The rest use run_sql; 12% of those need a second try after a SQL error, which the agent handles on its own.

Follow-up questions to expect

• "How do you stop SQL injection?" — Bind parameters in fixed queries, and for generated SQL rely on a read-only role and allowlisted tables, not string checks.
• "The model keeps using wrong column names. What do you do?" — Put the schema and sample rows in the prompt, add column descriptions, and return SQL errors to the model.
• "How do you handle huge tables?" — Force a LIMIT, set a statement timeout, and prefer aggregate queries; summarise results before returning them.