github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-09 23:12:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Harrison/align table (#1081)

Browse Source 
Co-authored-by: Francisco Ingham <fpingham@gmail.com>
This commit is contained in:
Harrison Chase
2023-02-15 23:53:37 -08:00
committed by GitHub
parent c60954d0f8
commit 5e10e19bfe
4 changed files with 137 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
tests/unit_tests/test_sql_database.py
View File

@@ -28,12 +28,28 @@ def test_table_info() -> None:
metadata_obj.create_all(engine)
db = SQLDatabase(engine)
output = db.table_info
output = output[len(_TEMPLATE_PREFIX) :]
expected_output = (
"Table 'user' has columns: {'user_id': ['INTEGER'], 'user_name': ['VARCHAR(16)']}",
"Table 'company' has columns: {'company_id': ['INTEGER'], 'company_location': ['VARCHAR']}",
expected_output = """
CREATE TABLE user (
user_id INTEGER NOT NULL,
user_name VARCHAR(16) NOT NULL,
PRIMARY KEY (user_id)
)
assert sorted(output.split("\n")) == sorted(expected_output)
SELECT * FROM 'user' LIMIT 3
user_id user_name
CREATE TABLE company (
company_id INTEGER NOT NULL,
company_location VARCHAR NOT NULL,
PRIMARY KEY (company_id)
)
SELECT * FROM 'company' LIMIT 3
company_id company_location
"""
assert sorted(" ".join(output.split())) == sorted(" ".join(expected_output.split()))
def test_table_info_w_sample_rows() -> None:
@@ -51,12 +67,31 @@ def test_table_info_w_sample_rows() -> None:
db = SQLDatabase(engine, sample_rows_in_table_info=2)
output = db.table_info
output = output[len(_TEMPLATE_PREFIX) :]
expected_output = (
"Table 'user' has columns: {'user_id': ['INTEGER', ['13', '14']], 'user_name': ['VARCHAR(16)', ['Harrison', 'Chase']]}",
"Table 'company' has columns: {'company_id': ['INTEGER', []], 'company_location': ['VARCHAR', []]}",
)
assert sorted(output.split("\n")) == sorted(expected_output)
expected_output = """
CREATE TABLE company (
company_id INTEGER NOT NULL,
company_location VARCHAR NOT NULL,
PRIMARY KEY (company_id)
)
SELECT * FROM 'company' LIMIT 2
company_id company_location
CREATE TABLE user (
user_id INTEGER NOT NULL,
user_name VARCHAR(16) NOT NULL,
PRIMARY KEY (user_id)
)
SELECT * FROM 'user' LIMIT 2
user_id user_name
13 Harrison
14 Chase
"""
assert sorted(output.split()) == sorted(expected_output.split())
def test_sql_database_run() -> None:

16
tests/unit_tests/test_sql_database_schema.py
View File

@@ -1,3 +1,4 @@
# flake8: noqa
"""Test SQL database wrapper with schema support.
Using DuckDB as SQLite does not support schemas.
@@ -16,7 +17,7 @@ from sqlalchemy import (
schema,
)
from langchain.sql_database import _TEMPLATE_PREFIX, SQLDatabase
from langchain.sql_database import SQLDatabase
metadata_obj = MetaData()
@@ -46,11 +47,14 @@ def test_table_info() -> None:
metadata_obj.create_all(engine)
db = SQLDatabase(engine, schema="schema_a")
output = db.table_info
output = output[len(_TEMPLATE_PREFIX) :]
expected_output = (
"Table 'user' has columns: {'user_id': ['INTEGER'], 'user_name': ['VARCHAR']}"
)
assert output == expected_output
expected_output = """
CREATE TABLE schema_a."user"(user_id INTEGER, user_name VARCHAR NOT NULL, PRIMARY KEY(user_id));
SELECT * FROM 'user' LIMIT 3
user_id user_name
"""
assert sorted(" ".join(output.split())) == sorted(" ".join(expected_output.split()))
def test_sql_database_run() -> None: