From d17dea30ce1e3e412c65ac82754aba4f7b989432 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Tue, 4 Apr 2023 06:48:45 -0700 Subject: [PATCH] Harrison/sql views (#2376) Co-authored-by: Wadih Pazos Co-authored-by: Wadih Pazos Sr --- langchain/sql_database.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/langchain/sql_database.py b/langchain/sql_database.py index 234395a2..4193e619 100644 --- a/langchain/sql_database.py +++ b/langchain/sql_database.py @@ -22,6 +22,7 @@ class SQLDatabase: include_tables: Optional[List[str]] = None, sample_rows_in_table_info: int = 3, custom_table_info: Optional[dict] = None, + view_support: Optional[bool] = False, ): """Create engine from database URI.""" self._engine = engine @@ -30,7 +31,14 @@ class SQLDatabase: raise ValueError("Cannot specify both include_tables and ignore_tables") self._inspector = inspect(self._engine) - self._all_tables = set(self._inspector.get_table_names(schema=schema)) + + # including view support by adding the views as well as tables to the all + # tables list if view_support is True + self._all_tables = set( + self._inspector.get_table_names(schema=schema) + + (self._inspector.get_view_names(schema=schema) if view_support else []) + ) + self._include_tables = set(include_tables) if include_tables else set() if self._include_tables: missing_tables = self._include_tables - self._all_tables @@ -69,8 +77,12 @@ class SQLDatabase: ) self._metadata = metadata or MetaData() + # including view support if view_support = true self._metadata.reflect( - bind=self._engine, only=self._usable_tables, schema=self._schema + views=view_support, + bind=self._engine, + only=self._usable_tables, + schema=self._schema, ) @classmethod