harlequin-postgres 1.2.0__tar.gz → 1.2.2__tar.gz

{harlequin_postgres-1.2.0 → harlequin_postgres-1.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: harlequin-postgres
-Version: 1.2.0
+Version: 1.2.2
 Summary: A Harlequin adapter for Postgres.
 License: MIT
 Author: Ted Conbeer

{harlequin_postgres-1.2.0 → harlequin_postgres-1.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "harlequin-postgres"
-version = "1.2.0"
+version = "1.2.2"
 description = "A Harlequin adapter for Postgres."
 authors = ["Ted Conbeer <tconbeer@users.noreply.github.com>"]
 license = "MIT"

{harlequin_postgres-1.2.0 → harlequin_postgres-1.2.2}/src/harlequin_postgres/adapter.py

@@ -136,10 +136,19 @@ class HarlequinPostgresConnection(HarlequinConnection):
             cur.close()
             return None
         except Exception as e:
-            cur.close()
-            self.rollback()
+            msg_suffix = ""
+            try:
+                cur.close()
+                self.rollback()
+            except Exception:
+                # likely connection is closed; error messages
+                # can be cryptic, so help the user.
+                msg_suffix = (
+                    "\n\nYou may need to restart Harlequin to reconnect to the "
+                    "database."
+                )
             raise HarlequinQueryError(
-                msg=str(e),
+                msg=f"{e}{msg_suffix}",
                 title="Harlequin encountered an error while executing your query.",
             ) from e
         else:
@@ -217,15 +226,16 @@ class HarlequinPostgresConnection(HarlequinConnection):
         conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
-                f"""
+                """
                 select schema_name
                 from information_schema.schemata
                 where
-                    catalog_name = '{dbname}'
+                    catalog_name = %s
                     and schema_name != 'information_schema'
-                    and schema_name not like 'pg_%'
+                    and schema_name not like 'pg_%%'
                 order by schema_name asc
-                ;"""
+                ;""",
+                (dbname,),
             )
             results: list[tuple[str]] = cur.fetchall()
         self.pool.putconn(conn)
@@ -235,34 +245,78 @@ class HarlequinPostgresConnection(HarlequinConnection):
         conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
-                f"""
+                """
                 select table_name, table_type
                 from information_schema.tables
                 where
-                    table_catalog = '{dbname}'
-                    and table_schema = '{schema}'
+                    table_catalog = %s
+                    and table_schema = %s
                 order by table_name asc
-                ;"""
+                ;""",
+                (dbname, schema),
             )
             results: list[tuple[str, str]] = cur.fetchall()
         self.pool.putconn(conn)
         return results
 
+    # only works for the currently-connected db
+    def _get_mvs(self, schema: str) -> list[tuple[str]]:
+        conn: Connection = self.pool.getconn()
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                select matviewname
+                from pg_matviews
+                where
+                    schemaname = %s
+                order by matviewname asc
+                ;""",
+                (schema,),
+            )
+            results: list[tuple[str]] = cur.fetchall()
+        self.pool.putconn(conn)
+        return results
+
     def _get_columns(
         self, dbname: str, schema: str, relation: str
     ) -> list[tuple[str, str]]:
         conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
-                f"""
+                """
                 select column_name, data_type
                 from information_schema.columns
                 where
-                    table_catalog = '{dbname}'
-                    and table_schema = '{schema}'
-                    and table_name = '{relation}'
+                    table_catalog = %s
+                    and table_schema = %s
+                    and table_name = %s
                 order by ordinal_position asc
-                ;"""
+                ;""",
+                (dbname, schema, relation),
+            )
+            results: list[tuple[str, str]] = cur.fetchall()
+        self.pool.putconn(conn)
+        return results
+
+    def _get_mv_cols(self, schema: str, mv: str) -> list[tuple[str, str]]:
+        conn: Connection = self.pool.getconn()
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                select 
+                    a.attname,
+                    pg_catalog.format_type(a.atttypid, a.atttypmod)
+                from pg_attribute a
+                join pg_class t on a.attrelid = t.oid
+                join pg_namespace s on t.relnamespace = s.oid
+                where 
+                    a.attnum > 0 
+                    and not a.attisdropped
+                    and s.nspname = %s
+                    and t.relname = %s
+                order by a.attnum;
+                ;""",
+                (schema, mv),
             )
             results: list[tuple[str, str]] = cur.fetchall()
         self.pool.putconn(conn)

{harlequin_postgres-1.2.0 → harlequin_postgres-1.2.2}/src/harlequin_postgres/catalog.py

@@ -166,6 +166,38 @@ class ForeignCatalogItem(TableCatalogItem):
         )
 
 
+class MaterializedViewCatalogItem(RelationCatalogItem):
+    @classmethod
+    def from_parent(
+        cls,
+        parent: "SchemaCatalogItem",
+        label: str,
+    ) -> "MaterializedViewCatalogItem":
+        relation_query_name = f'"{parent.label}"."{label}"'
+        relation_qualified_identifier = f'{parent.qualified_identifier}."{label}"'
+        return cls(
+            qualified_identifier=relation_qualified_identifier,
+            query_name=relation_query_name,
+            label=label,
+            type_label="mv",
+            connection=parent.connection,
+            parent=parent,
+        )
+
+    def fetch_children(self) -> list[ColumnCatalogItem]:
+        if self.parent is None or self.parent.parent is None or self.connection is None:
+            return []
+        result = self.connection._get_mv_cols(self.parent.label, self.label)
+        return [
+            ColumnCatalogItem.from_parent(
+                parent=self,
+                label=column_name,
+                type_label=self.connection._short_column_type(column_type),
+            )
+            for column_name, column_type in result
+        ]
+
+
 @dataclass
 class SchemaCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]):
     INTERACTIONS = [
@@ -227,6 +259,14 @@ class SchemaCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]):
                     )
                 )
 
+        for (mv_label,) in self.connection._get_mvs(self.label):
+            children.append(
+                MaterializedViewCatalogItem.from_parent(
+                    parent=self,
+                    label=mv_label,
+                )
+            )
+
         return children