From 45fb40ae9578320beceac99fb03f5d6d03ed3a76 Mon Sep 17 00:00:00 2001 From: Helen Cristina Date: Sat, 11 Nov 2023 20:25:21 -0300 Subject: [PATCH] fix: Adjust `find-tables` to properly get Oracle and Postgres schemas (#1034) * Debugging find-tables issue * Remove debugging printing of schemas and table maps * fix: Add override method to list schemas for Postgres --- data_validation/__main__.py | 1 + data_validation/clients.py | 2 +- third_party/ibis/ibis_postgres/client.py | 13 +++++++++++++ 3 files changed, 15 insertions(+), 1 deletion(-) diff --git a/data_validation/__main__.py b/data_validation/__main__.py index 35e69df89..284fbf4c9 100644 --- a/data_validation/__main__.py +++ b/data_validation/__main__.py @@ -368,6 +368,7 @@ def get_table_map(client, allowed_schemas=None): """Return dict with searchable keys for table matching.""" table_map = {} table_objs = clients.get_all_tables(client, allowed_schemas=allowed_schemas) + for table_obj in table_objs: table_key = ".".join([t for t in table_obj if t]) table_map[table_key] = { diff --git a/data_validation/clients.py b/data_validation/clients.py index 3980b1832..262197619 100644 --- a/data_validation/clients.py +++ b/data_validation/clients.py @@ -182,7 +182,7 @@ def list_schemas(client): def list_tables(client, schema_name): """Return a list of tables in the DB schema.""" - if client.name in ["oracle", "postgres", "db2", "mssql", "redshift", "snowflake"]: + if client.name in ["db2", "mssql", "redshift", "snowflake"]: return client.list_tables() return client.list_tables(database=schema_name) diff --git a/third_party/ibis/ibis_postgres/client.py b/third_party/ibis/ibis_postgres/client.py index efd736b1e..09fa67f92 100644 --- a/third_party/ibis/ibis_postgres/client.py +++ b/third_party/ibis/ibis_postgres/client.py @@ -58,4 +58,17 @@ def _get_type(typestr: str) -> dt.DataType: return _parse_numeric(typestr) +def list_schemas(self, like=None): + with self.begin() as con: + # Databases on Postgres are not the same as schemas and for this method we need the schema list (SQL query reference: https://dba.stackexchange.com/a/127668) + schemas = [ + row.nspname + for row in con.exec_driver_sql( + "SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname !~ '^pg_' AND nspname <> 'information_schema' ORDER BY 1" + ).mappings() + ] + return self._filter_with_like(schemas, like) + + PostgresBackend._metadata = _metadata +PostgresBackend.list_databases = list_schemas