feat(postgres): support materialized views in get model TCTC-7444 (#1404

) * feat(postgres): support materialized views in get model Signed-off-by: Luka Peschke <[email protected]> * fix query syntax Signed-off-by: Luka Peschke <[email protected]> * update changelog Signed-off-by: Luka Peschke <[email protected]> * feat: make query more cursed Signed-off-by: Luka Peschke <[email protected]> * tests: add materialized view Signed-off-by: Luka Peschke <[email protected]> * fix GROUP BY clause Signed-off-by: Luka Peschke <[email protected]> * feat: fix type names Signed-off-by: Luka Peschke <[email protected]> --------- Signed-off-by: Luka Peschke <[email protected]>
ToucanToco · Dec 15, 2023 · 3addaf7 · 3addaf7
1 parent b3bbe5b
commit 3addaf7
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog (Pypi package)
 
+## Unreleased
+
+### Changed
+
+- Postgres: Materialized views are now returned as well via `get_model`. Their type is `'view'`.
+
 ## [4.9.6] 2023-11-23
 
 ### Fixed

diff --git a/tests/postgres/fixtures/world_postgres.sql b/tests/postgres/fixtures/world_postgres.sql
@@ -5399,3 +5399,5 @@ COMMIT;
 ANALYZE city;
 ANALYZE country;
 ANALYZE countrylanguage;
+
+CREATE MATERIALIZED VIEW country_materialized_view AS SELECT * FROM country;
diff --git a/tests/postgres/test_postgres.py b/tests/postgres/test_postgres.py
@@ -110,6 +110,35 @@ def postgres_db_model() -> list[dict]:
     ]
 
 
+@pytest.fixture
+def postgres_db_model_with_materialized_views(postgres_db_model: list[dict]) -> list[dict]:
+    return postgres_db_model + [
+        {
+            'database': 'postgres_db',
+            'name': 'country_materialized_view',
+            'schema': 'public',
+            'type': 'view',
+            'columns': [
+                {'name': 'code', 'type': 'character'},
+                {'name': 'name', 'type': 'text'},
+                {'name': 'continent', 'type': 'text'},
+                {'name': 'region', 'type': 'text'},
+                {'name': 'surfacearea', 'type': 'real'},
+                {'name': 'indepyear', 'type': 'smallint'},
+                {'name': 'population', 'type': 'integer'},
+                {'name': 'lifeexpectancy', 'type': 'real'},
+                {'name': 'gnp', 'type': 'numeric'},
+                {'name': 'gnpold', 'type': 'numeric'},
+                {'name': 'localname', 'type': 'text'},
+                {'name': 'governmentform', 'type': 'text'},
+                {'name': 'headofstate', 'type': 'text'},
+                {'name': 'capital', 'type': 'integer'},
+                {'name': 'code2', 'type': 'character'},
+            ],
+        }
+    ]
+
+
 def test_get_status_all_good(postgres_connector):
     assert postgres_connector.get_status() == ConnectorStatus(
         status=True,
@@ -454,6 +483,19 @@ def test_get_model(postgres_connector: PostgresConnector, postgres_db_model: lis
     assert postgres_connector.get_model(db_name='another_db') == []
 
 
+def test_get_model_with_materialized_views(
+    postgres_connector: PostgresConnector, postgres_db_model_with_materialized_views: list[dict]
+) -> None:
+    """Check that it returns the db tree structure"""
+    postgres_connector.include_materialized_views = True
+    assert postgres_connector.get_model() == postgres_db_model_with_materialized_views
+    assert (
+        postgres_connector.get_model(db_name='postgres_db')
+        == postgres_db_model_with_materialized_views
+    )
+    assert postgres_connector.get_model(db_name='another_db') == []
+
+
 def test_raised_error_for_get_model(mocker, postgres_connector):
     """Check that it returns the db tree structure"""
     with mocker.patch.object(

diff --git a/toucan_connectors/postgres/postgresql_connector.py b/toucan_connectors/postgres/postgresql_connector.py
@@ -110,6 +110,10 @@ class PostgresConnector(ToucanConnector, DiscoverableConnector, VersionableEngin
         'time you want to wait for the server to respond. None by default',
     )
 
+    include_materialized_views: bool = Field(
+        False, description='Wether materialized views should be listed in the query builder or not.'
+    )
+
     def get_connection_params(self, *, database: str | None = None):
         con_params = dict(
             user=self.user,
@@ -238,7 +242,11 @@ def _list_tables_info(self, database_name: str = None) -> List[tuple]:
             )
         )
         with connection.cursor() as cursor:
-            cursor.execute(build_database_model_extraction_query())
+            cursor.execute(
+                build_database_model_extraction_query(
+                    database_name, include_materialized_views=self.include_materialized_views
+                )
+            )
             return cursor.fetchall()
 
     def get_engine_version(self) -> tuple:

diff --git a/toucan_connectors/postgres/utils.py b/toucan_connectors/postgres/utils.py
@@ -603,7 +603,40 @@
 }
 
 
-def build_database_model_extraction_query() -> str:
+def _build_materialized_views_info_extraction_query(db_name: str | None) -> str:
+    # Here, we need to query pg_catalog because materialized views are not a standard SQL feature
+    # and are thus not available in information_schema.
+    # The WHERE condition filters to retrieve materialized views only and excludes system columns
+    # (see https://www.postgresql.org/docs/current/catalog-pg-attribute.html).
+    # The CASE statement is to format postgres types to sql types as found in information_schema
+    database_name = f"'{db_name}'" if db_name else 'NULL'
+    return f"""
+    SELECT {database_name} AS "database",
+    ns.nspname AS "schema",
+    'view' AS "table_type",
+    cls.relname AS table_name,
+    JSON_AGG(
+        JSON_BUILD_OBJECT(
+            'name', attr.attname,
+            'type', CASE
+                WHEN tp.typname = 'int2' THEN 'smallint'
+                WHEN tp.typname = 'bpchar' THEN 'character'
+                WHEN tp.typname IN ('int4', 'int8') THEN 'integer'
+                WHEN tp.typname IN ('float2', 'float4', 'float8') THEN 'real'
+            ELSE tp.typname
+            END
+        )
+    ) AS columns
+    FROM pg_catalog.pg_attribute AS attr
+    JOIN pg_catalog.pg_class AS cls ON cls.oid = attr.attrelid
+    JOIN pg_catalog.pg_namespace AS ns ON ns.oid = cls.relnamespace
+    JOIN pg_catalog.pg_type AS tp ON tp.oid = attr.atttypid
+    WHERE cls.relname in (SELECT matviewname FROM pg_matviews) AND attr.attnum >= 0
+    GROUP BY schema, database, table_name, table_type
+    """
+
+
+def _build_regular_tables_model_extraction_query() -> str:
     return """SELECT t.table_catalog AS database,
     t.table_schema AS schema,
     CASE WHEN t.table_type = 'BASE TABLE' THEN 'table' ELSE lower(t.table_type) END AS type,
@@ -615,5 +648,18 @@ def build_database_model_extraction_query() -> str:
         t.table_name = c.table_name AND t.table_schema = c.table_schema
     WHERE t.table_type IN ('BASE TABLE', 'VIEW')
     AND t.table_schema NOT IN  ('pg_catalog', 'information_schema', 'pg_internal')
-    GROUP BY t.table_schema, t.table_catalog, t.table_name, t.table_type;
+    GROUP BY t.table_schema, t.table_catalog, t.table_name, t.table_type
     """
+
+
+def build_database_model_extraction_query(
+    db_name: str | None, include_materialized_views: bool
+) -> str:
+    return (
+        f"""{_build_regular_tables_model_extraction_query()}
+        UNION ALL
+        {_build_materialized_views_info_extraction_query(db_name)};
+        """
+        if include_materialized_views
+        else _build_regular_tables_model_extraction_query() + ';'
+    )