Harrison/sql views (#2376)

Co-authored-by: Wadih Pazos <wadih@wpazos.com>
Co-authored-by: Wadih Pazos Sr <wadih@esgenio.com>
This commit is contained in:
Harrison Chase 2023-04-04 06:48:45 -07:00 committed by GitHub
parent e90d007db3
commit d17dea30ce
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -22,6 +22,7 @@ class SQLDatabase:
include_tables: Optional[List[str]] = None, include_tables: Optional[List[str]] = None,
sample_rows_in_table_info: int = 3, sample_rows_in_table_info: int = 3,
custom_table_info: Optional[dict] = None, custom_table_info: Optional[dict] = None,
view_support: Optional[bool] = False,
): ):
"""Create engine from database URI.""" """Create engine from database URI."""
self._engine = engine self._engine = engine
@ -30,7 +31,14 @@ class SQLDatabase:
raise ValueError("Cannot specify both include_tables and ignore_tables") raise ValueError("Cannot specify both include_tables and ignore_tables")
self._inspector = inspect(self._engine) 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() self._include_tables = set(include_tables) if include_tables else set()
if self._include_tables: if self._include_tables:
missing_tables = self._include_tables - self._all_tables missing_tables = self._include_tables - self._all_tables
@ -69,8 +77,12 @@ class SQLDatabase:
) )
self._metadata = metadata or MetaData() self._metadata = metadata or MetaData()
# including view support if view_support = true
self._metadata.reflect( 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 @classmethod