diff --git a/docs/docs/integrations/tools/sql_database.ipynb b/docs/docs/integrations/tools/sql_database.ipynb new file mode 100644 index 00000000000..eec9027e605 --- /dev/null +++ b/docs/docs/integrations/tools/sql_database.ipynb @@ -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": [ + "\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": [ + "\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": [ + "\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 +} diff --git a/libs/community/langchain_community/tools/sql_database/tool.py b/libs/community/langchain_community/tools/sql_database/tool.py index 850cecb4510..4c7fe4726a0 100644 --- a/libs/community/langchain_community/tools/sql_database/tool.py +++ b/libs/community/langchain_community/tools/sql_database/tool.py @@ -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) diff --git a/libs/community/langchain_community/utilities/sql_database.py b/libs/community/langchain_community/utilities/sql_database.py index 298b8ca7e22..087ac44bc11 100644 --- a/libs/community/langchain_community/utilities/sql_database.py +++ b/libs/community/langchain_community/utilities/sql_database.py @@ -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}" diff --git a/libs/experimental/langchain_experimental/sql/base.py b/libs/experimental/langchain_experimental/sql/base.py index 1785598e36a..b6de5e67c9d 100644 --- a/libs/experimental/langchain_experimental/sql/base.py +++ b/libs/experimental/langchain_experimental/sql/base.py @@ -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 diff --git a/libs/experimental/langchain_experimental/sql/vector_sql.py b/libs/experimental/langchain_experimental/sql/vector_sql.py index 1b0db959044..a21a6bae868 100644 --- a/libs/experimental/langchain_experimental/sql/vector_sql.py +++ b/libs/experimental/langchain_experimental/sql/vector_sql.py @@ -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 diff --git a/libs/langchain/tests/unit_tests/test_sql_database.py b/libs/langchain/tests/unit_tests/test_sql_database.py index fea54e74c25..c593faae8e8 100644 --- a/libs/langchain/tests/unit_tests/test_sql_database.py +++ b/libs/langchain/tests/unit_tests/test_sql_database.py @@ -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"