github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-08-31 10:23:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[minor]: SQLDatabase Add fetch mode cursor, query parameters, query by selectable, expose execution options, and documentation (#17191)

Browse Source 
- **Description:** Improve `SQLDatabase` adapter component to promote
code re-use, see
[suggestion](https://github.com/langchain-ai/langchain/pull/16246#pullrequestreview-1846590962).
  - **Needed by:** GH-16246
  - **Addressed to:** @baskaryan, @cbornet 

## Details
- Add `cursor` fetch mode
- Accept SQL query parameters
- Accept both `str` and SQLAlchemy selectables as query expression
- Expose `execution_options`
- Documentation page (notebook) about `SQLDatabase` [^1]
See [About
SQLDatabase](https://github.com/langchain-ai/langchain/blob/c1c7b763/docs/docs/integrations/tools/sql_database.ipynb).

[^1]: Apparently there hasn't been any yet?

---------

Co-authored-by: Andreas Motl <andreas.motl@crate.io>
This commit is contained in:
Eugene Yurtsev
2024-02-07 22:23:43 -05:00
committed by GitHub
parent 7e4b676d53
commit 780e84ae79
6 changed files with 600 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

440
docs/docs/integrations/tools/sql_database.ipynb Normal file
View File

@@ -0,0 +1,440 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"# SQL Database\n",
"\n",
"::: {.callout-note}\n",
"The `SQLDatabase` adapter utility is a wrapper around a database connection.\n",
"\n",
"For talking to SQL databases, it uses the [SQLAlchemy] Core API .\n",
":::\n",
"\n",
"\n",
"This notebook shows how to use the utility to access an SQLite database.\n",
"It uses the example [Chinook Database], and demonstrates those features:\n",
"\n",
"- Query using SQL\n",
"- Query using SQLAlchemy selectable\n",
"- Fetch modes `cursor`, `all`, and `one`\n",
"- Bind query parameters\n",
"\n",
"[Chinook Database]: https://github.com/lerocha/chinook-database\n",
"[SQLAlchemy]: https://www.sqlalchemy.org/\n",
"\n",
"\n",
"You can use the `Tool` or `@tool` decorator to create a tool from this utility.\n",
"\n",
"\n",
"::: {.callout-caution}\n",
"If creating a tool from the SQLDatbase utility and combining it with an LLM or exposing it to an end user\n",
"remember to follow good security practices.\n",
"\n",
"See security information: https://python.langchain.com/docs/security\n",
":::"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true,
"jupyter": {
"outputs_hidden": true
}
},
"outputs": [],
"source": [
"!wget 'https://github.com/lerocha/chinook-database/releases/download/v1.4.2/Chinook_Sqlite.sql'"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1|AC/DC\r\n",
"2|Accept\r\n",
"3|Aerosmith\r\n",
"4|Alanis Morissette\r\n",
"5|Alice In Chains\r\n",
"6|Antônio Carlos Jobim\r\n",
"7|Apocalyptica\r\n",
"8|Audioslave\r\n",
"9|BackBeat\r\n",
"10|Billy Cobham\r\n",
"11|Black Label Society\r\n",
"12|Black Sabbath\r\n"
]
}
],
"source": [
"!sqlite3 -bail -cmd '.read Chinook_Sqlite.sql' -cmd 'SELECT * FROM Artist LIMIT 12;' -cmd '.quit'"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [],
"source": [
"!sqlite3 -bail -cmd '.read Chinook_Sqlite.sql' -cmd '.save Chinook.db' -cmd '.quit'"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Initialize Database"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [],
"source": [
"from pprint import pprint\n",
"\n",
"import sqlalchemy as sa\n",
"from langchain.sql_database import SQLDatabase\n",
"\n",
"db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query as cursor\n",
"\n",
"The fetch mode `cursor` returns results as SQLAlchemy's\n",
"`CursorResult` instance."
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<class 'sqlalchemy.engine.cursor.CursorResult'>\n",
"[{'ArtistId': 1, 'Name': 'AC/DC'},\n",
" {'ArtistId': 2, 'Name': 'Accept'},\n",
" {'ArtistId': 3, 'Name': 'Aerosmith'},\n",
" {'ArtistId': 4, 'Name': 'Alanis Morissette'},\n",
" {'ArtistId': 5, 'Name': 'Alice In Chains'},\n",
" {'ArtistId': 6, 'Name': 'Antônio Carlos Jobim'},\n",
" {'ArtistId': 7, 'Name': 'Apocalyptica'},\n",
" {'ArtistId': 8, 'Name': 'Audioslave'},\n",
" {'ArtistId': 9, 'Name': 'BackBeat'},\n",
" {'ArtistId': 10, 'Name': 'Billy Cobham'},\n",
" {'ArtistId': 11, 'Name': 'Black Label Society'},\n",
" {'ArtistId': 12, 'Name': 'Black Sabbath'}]\n"
]
}
],
"source": [
"result = db.run(\"SELECT * FROM Artist LIMIT 12;\", fetch=\"cursor\")\n",
"print(type(result))\n",
"pprint(list(result.mappings()))"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query as string payload\n",
"\n",
"The fetch modes `all` and `one` return results in string format."
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<class 'str'>\n",
"[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham'), (11, 'Black Label Society'), (12, 'Black Sabbath')]\n"
]
}
],
"source": [
"result = db.run(\"SELECT * FROM Artist LIMIT 12;\", fetch=\"all\")\n",
"print(type(result))\n",
"print(result)"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<class 'str'>\n",
"[(1, 'AC/DC')]\n"
]
}
],
"source": [
"result = db.run(\"SELECT * FROM Artist LIMIT 12;\", fetch=\"one\")\n",
"print(type(result))\n",
"print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query with parameters\n",
"\n",
"In order to bind query parameters, use the optional `parameters` argument."
]
},
{
"cell_type": "code",
"execution_count": 41,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'ArtistId': 35, 'Name': 'Pedro Luís & A Parede'},\n",
" {'ArtistId': 115, 'Name': 'Page & Plant'},\n",
" {'ArtistId': 116, 'Name': 'Passengers'},\n",
" {'ArtistId': 117, 'Name': \"Paul D'Ianno\"},\n",
" {'ArtistId': 118, 'Name': 'Pearl Jam'},\n",
" {'ArtistId': 119, 'Name': 'Peter Tosh'},\n",
" {'ArtistId': 120, 'Name': 'Pink Floyd'},\n",
" {'ArtistId': 121, 'Name': 'Planet Hemp'},\n",
" {'ArtistId': 186, 'Name': 'Pedro Luís E A Parede'},\n",
" {'ArtistId': 256, 'Name': 'Philharmonia Orchestra & Sir Neville Marriner'},\n",
" {'ArtistId': 275, 'Name': 'Philip Glass Ensemble'}]\n"
]
}
],
"source": [
"result = db.run(\n",
" \"SELECT * FROM Artist WHERE Name LIKE :search;\",\n",
" parameters={\"search\": \"p%\"},\n",
" fetch=\"cursor\",\n",
")\n",
"pprint(list(result.mappings()))"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query with SQLAlchemy selectable\n",
"\n",
"Other than plain-text SQL statements, the adapter also accepts SQLAlchemy selectables."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'ArtistId': 35, 'Name': 'Pedro Luís & A Parede'},\n",
" {'ArtistId': 115, 'Name': 'Page & Plant'},\n",
" {'ArtistId': 116, 'Name': 'Passengers'},\n",
" {'ArtistId': 117, 'Name': \"Paul D'Ianno\"},\n",
" {'ArtistId': 118, 'Name': 'Pearl Jam'},\n",
" {'ArtistId': 119, 'Name': 'Peter Tosh'},\n",
" {'ArtistId': 120, 'Name': 'Pink Floyd'},\n",
" {'ArtistId': 121, 'Name': 'Planet Hemp'},\n",
" {'ArtistId': 186, 'Name': 'Pedro Luís E A Parede'},\n",
" {'ArtistId': 256, 'Name': 'Philharmonia Orchestra & Sir Neville Marriner'},\n",
" {'ArtistId': 275, 'Name': 'Philip Glass Ensemble'}]\n"
]
}
],
"source": [
"# In order to build a selectable on SA's Core API, you need a table definition.\n",
"metadata = sa.MetaData()\n",
"artist = sa.Table(\n",
" \"Artist\",\n",
" metadata,\n",
" sa.Column(\"ArtistId\", sa.INTEGER, primary_key=True),\n",
" sa.Column(\"Name\", sa.TEXT),\n",
")\n",
"\n",
"# Build a selectable with the same semantics of the recent query.\n",
"query = sa.select(artist).where(artist.c.Name.like(\"p%\"))\n",
"result = db.run(query, fetch=\"cursor\")\n",
"pprint(list(result.mappings()))"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query with execution options\n",
"\n",
"It is possible to augment the statement invocation with custom execution options.\n",
"For example, when applying a schema name translation, subsequent statements will\n",
"fail, because they try to hit a non-existing table."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"ename": "OperationalError",
"evalue": "(sqlite3.OperationalError) no such table: bar.Artist\n[SQL: SELECT bar.\"Artist\".\"ArtistId\", bar.\"Artist\".\"Name\" \nFROM bar.\"Artist\" \nWHERE bar.\"Artist\".\"Name\" LIKE ?]\n[parameters: ('p%',)]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1969\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[0;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[1;32m 1968\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m evt_handled:\n\u001b[0;32m-> 1969\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1970\u001b[0m \u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[1;32m 1971\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1973\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_has_events \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_has_events:\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/default.py:922\u001b[0m, in \u001b[0;36mDefaultDialect.do_execute\u001b[0;34m(self, cursor, statement, parameters, context)\u001b[0m\n\u001b[1;32m 921\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdo_execute\u001b[39m(\u001b[38;5;28mself\u001b[39m, cursor, statement, parameters, context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 922\u001b[0m cursor\u001b[38;5;241m.\u001b[39mexecute(statement, parameters)\n",
"\u001b[0;31mOperationalError\u001b[0m: no such table: bar.Artist",
"\nThe above exception was the direct cause of the following exception:\n",
"\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[6], line 3\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Build a selectable with the same semantics of the recent query.\u001b[39;00m\n\u001b[1;32m 2\u001b[0m query \u001b[38;5;241m=\u001b[39m sa\u001b[38;5;241m.\u001b[39mselect(artist)\u001b[38;5;241m.\u001b[39mwhere(artist\u001b[38;5;241m.\u001b[39mc\u001b[38;5;241m.\u001b[39mName\u001b[38;5;241m.\u001b[39mlike(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mp\u001b[39m\u001b[38;5;124m%\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m----> 3\u001b[0m \u001b[43mdb\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfetch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mresult\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mschema_translate_map\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43m{\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbar\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/libs/community/langchain_community/utilities/sql_database.py:484\u001b[0m, in \u001b[0;36mSQLDatabase.run\u001b[0;34m(self, command, fetch, parameters, execution_options, include_columns)\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrun\u001b[39m(\n\u001b[1;32m 472\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 473\u001b[0m command: Union[\u001b[38;5;28mstr\u001b[39m, Executable],\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 477\u001b[0m include_columns: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m 478\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[\u001b[38;5;28mstr\u001b[39m, Result]:\n\u001b[1;32m 479\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Execute a SQL command and return a string representing the results.\u001b[39;00m\n\u001b[1;32m 480\u001b[0m \n\u001b[1;32m 481\u001b[0m \u001b[38;5;124;03m If the statement returns rows, a string of the results is returned.\u001b[39;00m\n\u001b[1;32m 482\u001b[0m \u001b[38;5;124;03m If the statement returns no rows, an empty string is returned.\u001b[39;00m\n\u001b[1;32m 483\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 484\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 485\u001b[0m \u001b[43m \u001b[49m\u001b[43mcommand\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfetch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mexecution_options\u001b[49m\n\u001b[1;32m 486\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 488\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m fetch \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mresult\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m 489\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/libs/community/langchain_community/utilities/sql_database.py:450\u001b[0m, in \u001b[0;36mSQLDatabase._execute\u001b[0;34m(self, command, fetch, parameters, execution_options)\u001b[0m\n\u001b[1;32m 448\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 449\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mQuery expression has unknown type: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(command)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 450\u001b[0m cursor \u001b[38;5;241m=\u001b[39m \u001b[43mconnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 451\u001b[0m \u001b[43m \u001b[49m\u001b[43mcommand\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 452\u001b[0m \u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 453\u001b[0m \u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mexecution_options\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 454\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 456\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cursor\u001b[38;5;241m.\u001b[39mreturns_rows:\n\u001b[1;32m 457\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m fetch \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mall\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1416\u001b[0m, in \u001b[0;36mConnection.execute\u001b[0;34m(self, statement, parameters, execution_options)\u001b[0m\n\u001b[1;32m 1414\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mObjectNotExecutableError(statement) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n\u001b[1;32m 1415\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1416\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1417\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1418\u001b[0m \u001b[43m \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1419\u001b[0m \u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mNO_OPTIONS\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1420\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/sql/elements.py:516\u001b[0m, in \u001b[0;36mClauseElement._execute_on_connection\u001b[0;34m(self, connection, distilled_params, execution_options)\u001b[0m\n\u001b[1;32m 514\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m TYPE_CHECKING:\n\u001b[1;32m 515\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m, Executable)\n\u001b[0;32m--> 516\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mconnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute_clauseelement\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 517\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdistilled_params\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\n\u001b[1;32m 518\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 519\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 520\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mObjectNotExecutableError(\u001b[38;5;28mself\u001b[39m)\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1639\u001b[0m, in \u001b[0;36mConnection._execute_clauseelement\u001b[0;34m(self, elem, distilled_parameters, execution_options)\u001b[0m\n\u001b[1;32m 1627\u001b[0m compiled_cache: Optional[CompiledCacheType] \u001b[38;5;241m=\u001b[39m execution_options\u001b[38;5;241m.\u001b[39mget(\n\u001b[1;32m 1628\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcompiled_cache\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_compiled_cache\n\u001b[1;32m 1629\u001b[0m )\n\u001b[1;32m 1631\u001b[0m compiled_sql, extracted_params, cache_hit \u001b[38;5;241m=\u001b[39m elem\u001b[38;5;241m.\u001b[39m_compile_w_cache(\n\u001b[1;32m 1632\u001b[0m dialect\u001b[38;5;241m=\u001b[39mdialect,\n\u001b[1;32m 1633\u001b[0m compiled_cache\u001b[38;5;241m=\u001b[39mcompiled_cache,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1637\u001b[0m linting\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdialect\u001b[38;5;241m.\u001b[39mcompiler_linting \u001b[38;5;241m|\u001b[39m compiler\u001b[38;5;241m.\u001b[39mWARN_LINTING,\n\u001b[1;32m 1638\u001b[0m )\n\u001b[0;32m-> 1639\u001b[0m ret \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute_context\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1640\u001b[0m \u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1641\u001b[0m \u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecution_ctx_cls\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_init_compiled\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1642\u001b[0m \u001b[43m \u001b[49m\u001b[43mcompiled_sql\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1643\u001b[0m \u001b[43m \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1644\u001b[0m \u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1645\u001b[0m \u001b[43m \u001b[49m\u001b[43mcompiled_sql\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1646\u001b[0m \u001b[43m \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1647\u001b[0m \u001b[43m \u001b[49m\u001b[43melem\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1648\u001b[0m \u001b[43m \u001b[49m\u001b[43mextracted_params\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1649\u001b[0m \u001b[43m \u001b[49m\u001b[43mcache_hit\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_hit\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1650\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1651\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_events:\n\u001b[1;32m 1652\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mafter_execute(\n\u001b[1;32m 1653\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 1654\u001b[0m elem,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1658\u001b[0m ret,\n\u001b[1;32m 1659\u001b[0m )\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1848\u001b[0m, in \u001b[0;36mConnection._execute_context\u001b[0;34m(self, dialect, constructor, statement, parameters, execution_options, *args, **kw)\u001b[0m\n\u001b[1;32m 1843\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exec_insertmany_context(\n\u001b[1;32m 1844\u001b[0m dialect,\n\u001b[1;32m 1845\u001b[0m context,\n\u001b[1;32m 1846\u001b[0m )\n\u001b[1;32m 1847\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1848\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_exec_single_context\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1849\u001b[0m \u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstatement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\n\u001b[1;32m 1850\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1988\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[0;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[1;32m 1985\u001b[0m result \u001b[38;5;241m=\u001b[39m context\u001b[38;5;241m.\u001b[39m_setup_result_proxy()\n\u001b[1;32m 1987\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m-> 1988\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_handle_dbapi_exception\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1989\u001b[0m \u001b[43m \u001b[49m\u001b[43me\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[1;32m 1990\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1992\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:2343\u001b[0m, in \u001b[0;36mConnection._handle_dbapi_exception\u001b[0;34m(self, e, statement, parameters, cursor, context, is_sub_exec)\u001b[0m\n\u001b[1;32m 2341\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m should_wrap:\n\u001b[1;32m 2342\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m sqlalchemy_exception \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 2343\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m sqlalchemy_exception\u001b[38;5;241m.\u001b[39mwith_traceback(exc_info[\u001b[38;5;241m2\u001b[39m]) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m 2344\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 2345\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_info[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1969\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[0;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[1;32m 1967\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m 1968\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m evt_handled:\n\u001b[0;32m-> 1969\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1970\u001b[0m \u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[1;32m 1971\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1973\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_has_events \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_has_events:\n\u001b[1;32m 1974\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mafter_cursor_execute(\n\u001b[1;32m 1975\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 1976\u001b[0m cursor,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1980\u001b[0m context\u001b[38;5;241m.\u001b[39mexecutemany,\n\u001b[1;32m 1981\u001b[0m )\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/default.py:922\u001b[0m, in \u001b[0;36mDefaultDialect.do_execute\u001b[0;34m(self, cursor, statement, parameters, context)\u001b[0m\n\u001b[1;32m 921\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdo_execute\u001b[39m(\u001b[38;5;28mself\u001b[39m, cursor, statement, parameters, context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 922\u001b[0m cursor\u001b[38;5;241m.\u001b[39mexecute(statement, parameters)\n",
"\u001b[0;31mOperationalError\u001b[0m: (sqlite3.OperationalError) no such table: bar.Artist\n[SQL: SELECT bar.\"Artist\".\"ArtistId\", bar.\"Artist\".\"Name\" \nFROM bar.\"Artist\" \nWHERE bar.\"Artist\".\"Name\" LIKE ?]\n[parameters: ('p%',)]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)"
]
}
],
"source": [
"query = sa.select(artist).where(artist.c.Name.like(\"p%\"))\n",
"db.run(query, fetch=\"cursor\", execution_options={\"schema_translate_map\": {None: \"bar\"}})"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.2"
}
},
"nbformat": 4,
"nbformat_minor": 4
}

6
libs/community/langchain_community/tools/sql_database/tool.py
View File

@@ -1,6 +1,8 @@
# flake8: noqa
"""Tools for interacting with a SQL database."""
from typing import Any, Dict, Optional, Type
from typing import Any, Dict, Optional, Sequence, Type, Union
from sqlalchemy import Result
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
@@ -42,7 +44,7 @@ class QuerySQLDataBaseTool(BaseSQLDatabaseTool, BaseTool):
self,
query: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> str:
) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
"""Execute the query, return the results or an error message."""
return self.db.run_no_throw(query)

102
libs/community/langchain_community/utilities/sql_database.py
View File

@@ -1,12 +1,21 @@
"""SQLAlchemy wrapper around a database."""
from __future__ import annotations
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Union
import sqlalchemy
from langchain_core._api import deprecated
from langchain_core.utils import get_from_env
from sqlalchemy import MetaData, Table, create_engine, inspect, select, text
from sqlalchemy import (
Executable,
MetaData,
Result,
Table,
create_engine,
inspect,
select,
text,
)
from sqlalchemy.engine import Engine
from sqlalchemy.exc import ProgrammingError, SQLAlchemyError
from sqlalchemy.schema import CreateTable
@@ -373,67 +382,113 @@ class SQLDatabase:
def _execute(
self,
command: str,
fetch: Literal["all", "one"] = "all",
) -> Sequence[Dict[str, Any]]:
command: Union[str, Executable],
fetch: Literal["all", "one", "cursor"] = "all",
*,
parameters: Optional[Dict[str, Any]] = None,
execution_options: Optional[Dict[str, Any]] = None,
) -> Union[Sequence[Dict[str, Any]], Result]:
"""
Executes SQL command through underlying engine.
If the statement returns no rows, an empty list is returned.
"""
parameters = parameters or {}
execution_options = execution_options or {}
with self._engine.begin() as connection: # type: Connection # type: ignore[name-defined]
if self._schema is not None:
if self.dialect == "snowflake":
connection.exec_driver_sql(
"ALTER SESSION SET search_path = %s", (self._schema,)
"ALTER SESSION SET search_path = %s",
(self._schema,),
execution_options=execution_options,
)
elif self.dialect == "bigquery":
connection.exec_driver_sql("SET @@dataset_id=?", (self._schema,))
connection.exec_driver_sql(
"SET @@dataset_id=?",
(self._schema,),
execution_options=execution_options,
)
elif self.dialect == "mssql":
pass
elif self.dialect == "trino":
connection.exec_driver_sql("USE ?", (self._schema,))
connection.exec_driver_sql(
"USE ?",
(self._schema,),
execution_options=execution_options,
)
elif self.dialect == "duckdb":
# Unclear which parameterized argument syntax duckdb supports.
# The docs for the duckdb client say they support multiple,
# but `duckdb_engine` seemed to struggle with all of them:
# https://github.com/Mause/duckdb_engine/issues/796
connection.exec_driver_sql(f"SET search_path TO {self._schema}")
connection.exec_driver_sql(
f"SET search_path TO {self._schema}",
execution_options=execution_options,
)
elif self.dialect == "oracle":
connection.exec_driver_sql(
f"ALTER SESSION SET CURRENT_SCHEMA = {self._schema}"
f"ALTER SESSION SET CURRENT_SCHEMA = {self._schema}",
execution_options=execution_options,
)
elif self.dialect == "sqlany":
# If anybody using Sybase SQL anywhere database then it should not
# go to else condition. It should be same as mssql.
pass
elif self.dialect == "postgresql": # postgresql
connection.exec_driver_sql("SET search_path TO %s", (self._schema,))
connection.exec_driver_sql(
"SET search_path TO %s",
(self._schema,),
execution_options=execution_options,
)
if isinstance(command, str):
command = text(command)
elif isinstance(command, Executable):
pass
else:
raise TypeError(f"Query expression has unknown type: {type(command)}")
cursor = connection.execute(
command,
parameters,
execution_options=execution_options,
)
cursor = connection.execute(text(command))
if cursor.returns_rows:
if fetch == "all":
result = [x._asdict() for x in cursor.fetchall()]
elif fetch == "one":
first_result = cursor.fetchone()
result = [] if first_result is None else [first_result._asdict()]
elif fetch == "cursor":
return cursor
else:
raise ValueError("Fetch parameter must be either 'one' or 'all'")
raise ValueError(
"Fetch parameter must be either 'one', 'all', or 'cursor'"
)
return result
return []
def run(
self,
command: str,
fetch: Literal["all", "one"] = "all",
command: Union[str, Executable],
fetch: Literal["all", "one", "cursor"] = "all",
include_columns: bool = False,
) -> str:
*,
parameters: Optional[Dict[str, Any]] = None,
execution_options: Optional[Dict[str, Any]] = None,
) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
"""Execute a SQL command and return a string representing the results.
If the statement returns rows, a string of the results is returned.
If the statement returns no rows, an empty string is returned.
"""
result = self._execute(command, fetch)
result = self._execute(
command, fetch, parameters=parameters, execution_options=execution_options
)
if fetch == "cursor":
return result
res = [
{
@@ -472,7 +527,10 @@ class SQLDatabase:
command: str,
fetch: Literal["all", "one"] = "all",
include_columns: bool = False,
) -> str:
*,
parameters: Optional[Dict[str, Any]] = None,
execution_options: Optional[Dict[str, Any]] = None,
) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
"""Execute a SQL command and return a string representing the results.
If the statement returns rows, a string of the results is returned.
@@ -481,7 +539,13 @@ class SQLDatabase:
If the statement throws an error, the error message is returned.
"""
try:
return self.run(command, fetch, include_columns)
return self.run(
command,
fetch,
parameters=parameters,
execution_options=execution_options,
include_columns=include_columns,
)
except SQLAlchemyError as e:
"""Format the error message"""
return f"Error: {e}"

2
libs/experimental/langchain_experimental/sql/base.py
View File

@@ -173,7 +173,7 @@ class SQLDatabaseChain(Chain):
sql_cmd = checked_sql_command
_run_manager.on_text("\nSQLResult: ", verbose=self.verbose)
_run_manager.on_text(result, color="yellow", verbose=self.verbose)
_run_manager.on_text(str(result), color="yellow", verbose=self.verbose)
# If return direct, we just set the final result equal to
# the result of the sql query result, otherwise try to get a human readable
# final answer

1
libs/experimental/langchain_experimental/sql/vector_sql.py
View File

@@ -78,6 +78,7 @@ class VectorSQLRetrieveAllOutputParser(VectorSQLOutputParser):
def get_result_from_sqldb(db: SQLDatabase, cmd: str) -> Sequence[Dict[str, Any]]:
result = db._execute(cmd, fetch="all")
assert isinstance(result, Sequence)
return result

75
libs/langchain/tests/unit_tests/test_sql_database.py
View File

@@ -1,16 +1,19 @@
# flake8: noqa=E501
# flake8: noqa: E501
"""Test SQL database wrapper."""
import pytest
import sqlalchemy as sa
from langchain_community.utilities.sql_database import SQLDatabase, truncate_word
from sqlalchemy import (
Column,
Integer,
MetaData,
Result,
String,
Table,
Text,
create_engine,
insert,
select,
)
metadata_obj = MetaData()
@@ -108,8 +111,8 @@ def test_table_info_w_sample_rows() -> None:
assert sorted(output.split()) == sorted(expected_output.split())
def test_sql_database_run() -> None:
"""Test that commands can be run successfully and returned in correct format."""
def test_sql_database_run_fetch_all() -> None:
"""Verify running SQL expressions returning results as strings."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(
@@ -131,6 +134,52 @@ def test_sql_database_run() -> None:
assert full_output == expected_full_output
def test_sql_database_run_fetch_result() -> None:
"""Verify running SQL expressions returning results as SQLAlchemy `Result` instances."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn:
conn.execute(stmt)
db = SQLDatabase(engine)
command = "select user_id, user_name, user_bio from user where user_id = 17"
result = db.run(command, fetch="cursor", include_columns=True)
expected = [{"user_id": 17, "user_name": "hwchase", "user_bio": None}]
assert isinstance(result, Result)
assert result.mappings().fetchall() == expected
def test_sql_database_run_with_parameters() -> None:
"""Verify running SQL expressions with query parameters."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn:
conn.execute(stmt)
db = SQLDatabase(engine)
command = "select user_id, user_name, user_bio from user where user_id = :user_id"
full_output = db.run(command, parameters={"user_id": 17}, include_columns=True)
expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]"
assert full_output == expected_full_output
def test_sql_database_run_sqlalchemy_selectable() -> None:
"""Verify running SQL expressions using SQLAlchemy selectable."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn:
conn.execute(stmt)
db = SQLDatabase(engine)
command = select(user).where(user.c.user_id == 17)
full_output = db.run(command, include_columns=True)
expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]"
assert full_output == expected_full_output
def test_sql_database_run_update() -> None:
"""Test commands which return no rows return an empty string."""
engine = create_engine("sqlite:///:memory:")
@@ -145,6 +194,24 @@ def test_sql_database_run_update() -> None:
assert output == expected_output
def test_sql_database_schema_translate_map() -> None:
"""Verify using statement-specific execution options."""
engine = create_engine("sqlite:///:memory:")
db = SQLDatabase(engine)
# Define query using SQLAlchemy selectable.
command = select(user).where(user.c.user_id == 17)
# Define statement-specific execution options.
execution_options = {"schema_translate_map": {None: "bar"}}
# Verify the schema translation is applied.
with pytest.raises(sa.exc.OperationalError) as ex:
db.run(command, execution_options=execution_options, fetch="cursor")
assert ex.match("no such table: bar.user")
def test_truncate_word() -> None:
assert truncate_word("Hello World", length=5) == "He..."
assert truncate_word("Hello World", length=0) == "Hello World"