Merging upstream version 1.2.0.

Signed-off-by: Daniel Baumann <daniel@debian.org>

src/harlequin_postgres/catalog.py

@@ -13,7 +13,13 @@ from harlequin_postgres.interactions import (
     execute_drop_view_statement,
     execute_use_statement,
     insert_columns_at_cursor,
+    show_describe_relation,
+    show_describe_table_constraints,
+    show_describe_table_indexes,
+    show_list_indexes,
+    show_list_objects,
     show_select_star,
+    show_view_definition,
 )
 
 if TYPE_CHECKING:
@@ -49,6 +55,7 @@ class RelationCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"])
     INTERACTIONS = [
         ("Insert Columns at Cursor", insert_columns_at_cursor),
         ("Preview Data", show_select_star),
+        ("Describe Relation (\\d+)", show_describe_relation),
     ]
     parent: "SchemaCatalogItem" | None = None
 
@@ -70,6 +77,7 @@ class RelationCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"])
 
 class ViewCatalogItem(RelationCatalogItem):
     INTERACTIONS = RelationCatalogItem.INTERACTIONS + [
+        ("Show View Definition", show_view_definition),
         ("Drop View", execute_drop_view_statement),
     ]
 
@@ -93,6 +101,8 @@ class ViewCatalogItem(RelationCatalogItem):
 
 class TableCatalogItem(RelationCatalogItem):
     INTERACTIONS = RelationCatalogItem.INTERACTIONS + [
+        ("Describe Indexes", show_describe_table_indexes),
+        ("Describe Constraints", show_describe_table_constraints),
         ("Drop Table", execute_drop_table_statement),
     ]
 
@@ -115,10 +125,6 @@ class TableCatalogItem(RelationCatalogItem):
 
 
 class TempTableCatalogItem(TableCatalogItem):
-    INTERACTIONS = RelationCatalogItem.INTERACTIONS + [
-        ("Drop Table", execute_drop_table_statement),
-    ]
-
     @classmethod
     def from_parent(
         cls,
@@ -164,6 +170,8 @@ class ForeignCatalogItem(TableCatalogItem):
 class SchemaCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]):
     INTERACTIONS = [
         ("Set Search Path", execute_use_statement),
+        ("List Relations (\\d+)", show_list_objects),
+        ("List Indexes (\\di+)", show_list_indexes),
         ("Drop Schema", execute_drop_schema_statement),
     ]
     parent: "DatabaseCatalogItem" | None = None
@@ -224,6 +232,8 @@ class SchemaCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]):
 
 class DatabaseCatalogItem(InteractiveCatalogItem["HarlequinPostgresConnection"]):
     INTERACTIONS = [
+        ("List Relations (\\d+)", show_list_objects),
+        ("List Indexes (\\di+)", show_list_indexes),
         ("Drop Database", execute_drop_database_statement),
     ]
 

src/harlequin_postgres/interactions.py

@@ -8,11 +8,13 @@ from harlequin.exception import HarlequinQueryError
 
 if TYPE_CHECKING:
     from harlequin.driver import HarlequinDriver
-    from harlequin_duckdb.catalog import (
+
+    from harlequin_postgres.catalog import (
         ColumnCatalogItem,
         DatabaseCatalogItem,
         RelationCatalogItem,
         SchemaCatalogItem,
+        ViewCatalogItem,
     )
 
 
@@ -132,6 +134,459 @@ def show_select_star(
     )
 
 
+def show_list_objects(
+    item: "SchemaCatalogItem" | "DatabaseCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    # sourced from psql with -E, then the following command:
+    # \dtvmsE+ <schema>.*
+
+    # can't use isinstance due to circular reference
+    if type(item).__name__ == "SchemaCatalogItem":
+        where_clause = f"and n.nspname = '{item.label}'"
+    else:
+        where_clause = (
+            "and n.nspname not in ('pg_catalog', 'pg_toast', 'information_schema')"
+        )
+    driver.insert_text_in_new_buffer(
+        dedent(
+            f"""
+            select 
+                n.nspname as "Schema",
+                c.relname as "Name",
+                case c.relkind 
+                    when 'r' then 'table'
+                    when 'v' then 'view' 
+                    when 'm' then 'materialized view'
+                    when 'S'
+                    then 'sequence'
+                    when 't' then 'TOAST table' 
+                    when 'f' then 'foreign table' 
+                    when 'p' then 'partitioned table' 
+                end as "Type",
+                pg_catalog.pg_get_userbyid(c.relowner) as "Owner",
+                case c.relpersistence 
+                    when 'p' then 'permanent' 
+                    when 't' then 'temporary' 
+                    when 'u' then 'unlogged' 
+                end as "Persistence",
+                am.amname as "Access method",
+                pg_catalog.pg_size_pretty(pg_catalog.pg_table_size(c.oid)) as "Size",
+                pg_catalog.obj_description(c.oid, 'pg_class') as "Description"
+            from pg_catalog.pg_class c
+                left join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+                left join pg_catalog.pg_am am on am.oid = c.relam
+            where
+                c.relkind IN ('r','p','t','v','m', 's', 'S', 'f')
+                {where_clause}
+            order by 1,2;
+            """.strip("\n")
+        )
+    )
+
+
+def show_list_indexes(
+    item: "SchemaCatalogItem" | "DatabaseCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    # sourced from psql with -E, then the following command:
+    # \dis+ <schema>.*
+
+    # can't use isinstance due to circular reference
+    if type(item).__name__ == "SchemaCatalogItem":
+        where_clause = f"and n.nspname = '{item.label}'"
+    else:
+        where_clause = (
+            "and n.nspname not in ('pg_catalog', 'pg_toast', 'information_schema')"
+        )
+    driver.insert_text_in_new_buffer(
+        dedent(
+            f"""
+            select
+                n.nspname as "Schema",
+                c.relname as "Name",
+                case
+                    c.relkind
+                    when 'i'
+                    then 'index'
+                    when 'I'
+                    then 'partitioned index'
+                end as "Type",
+                pg_catalog.pg_get_userbyid(c.relowner) as "Owner",
+                c2.relname as "Table",
+                case
+                    c.relpersistence
+                    when 'p'
+                    then 'permanent'
+                    when 't'
+                    then 'temporary'
+                    when 'u'
+                    then 'unlogged'
+                end as "Persistence",
+                am.amname as "Access method",
+                pg_catalog.pg_size_pretty(pg_catalog.pg_table_size(c.oid)) as "Size",
+                pg_catalog.obj_description(c.oid, 'pg_class') as "Description"
+            from pg_catalog.pg_class c
+            left join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+            left join pg_catalog.pg_am am on am.oid = c.relam
+            left join pg_catalog.pg_index i on i.indexrelid = c.oid
+            left join pg_catalog.pg_class c2 on i.indrelid = c2.oid
+            where
+                c.relkind in ('i', 'I')
+                {where_clause}
+            order by 1,2;
+            """.strip("\n")
+        )
+    )
+
+
+def show_describe_relation(
+    item: "RelationCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    # sourced from psql -E \d+ {my rel}
+    # see https://stackoverflow.com/questions/60155968/using-results-of-d-command-in-psql
+    if item.parent is None:
+        driver.notify(
+            f"Could not describe {item.label} due to missing schema reference.",
+            severity="error",
+        )
+        return
+    driver.insert_text_in_new_buffer(
+        dedent(
+            f"""
+            with
+                index_columns as (
+                    select i.indexrelid, c.oid as rel_oid, unnest(i.indkey) as attnum
+                    from pg_catalog.pg_index i
+                    join pg_catalog.pg_class c on c.oid = i.indrelid
+                    where c.relname = '{item.label}'
+                ),
+                index_column_counts as (
+                    select rel_oid, attnum, count(*) as cnt 
+                    from index_columns 
+                    group by 1, 2
+                ),
+                constraint_columns as (
+                    select con.oid, c.oid as rel_oid, unnest(con.conkey) as attnum
+                    from pg_catalog.pg_constraint con
+                    join pg_catalog.pg_class c on con.conrelid = c.oid
+                    where c.relname = '{item.label}'
+                ),
+                constraint_column_counts as (
+                    select rel_oid, attnum, count(*) as cnt
+                    from constraint_columns
+                    group by 1, 2
+                ),
+                fkey_columns as (
+                    select
+                        src.relname as src_name,
+                        src.relnamespace::regnamespace as src_schema,
+                        c.oid as rel_oid,
+                        unnest(con.confkey) as attnum
+                    from pg_catalog.pg_constraint con
+                    join pg_catalog.pg_class c on con.confrelid = c.oid
+                    join pg_catalog.pg_class src on con.conrelid = src.oid
+                    where c.relname = '{item.label}'
+                ),
+                fkey_references as (
+                    select
+                        rel_oid, 
+                        attnum, 
+                        string_agg(src_schema || '.' || src_name, ', ') as sources
+                    from fkey_columns
+                    group by 1, 2
+                )
+            select
+                a.attname as "Column",
+                pg_catalog.format_type(a.atttypid, a.atttypmod) as "Type",
+                coll.collname as "Collation",
+                case
+                    when a.attnotnull is true then 'not null' 
+                    else '' 
+                end as "Nullable",
+                pg_catalog.pg_get_expr(d.adbin, d.adrelid, true) as "Default",
+                case
+                    a.attstorage
+                    when 'p'
+                    then 'plain'
+                    when 'x'
+                    then 'extended'
+                    when 'e'
+                    then 'external'
+                    when 'm'
+                    then 'main'
+                    else a.attstorage::text
+                end as "Storage",
+                case
+                    a.attcompression
+                    when 'p'
+                    then 'pglz'
+                    when 'l'
+                    then 'LZ4'
+                    else a.attcompression::text
+                end as "Compression",
+                case 
+                    when a.attstattarget = -1
+                    then null
+                    else a.attstattarget
+                end as "Stats target",
+                case 
+                    when index_column_counts.cnt > 0 then true else false
+                end as "Has Index",
+                case
+                    when constraint_column_counts.cnt > 0 then true else false
+                end as "Has Constraint",
+                fkey_references.sources as "Referenced by",
+                pg_catalog.col_description(a.attrelid, a.attnum) as "Description"
+            from pg_catalog.pg_attribute a
+            join pg_catalog.pg_class c on a.attrelid = c.oid
+            left join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+            left join pg_catalog.pg_collation coll on coll.oid = a.attcollation
+            left join
+                pg_catalog.pg_type t
+                on (t.oid = a.atttypid and t.typcollation <> a.attcollation)
+            left join
+                pg_catalog.pg_attrdef d
+                on (a.attrelid = d.adrelid and a.attnum = d.adnum and a.atthasdef)
+            left join
+                index_column_counts
+                on a.attnum = index_column_counts.attnum
+                and a.attrelid = index_column_counts.rel_oid
+            left join
+                constraint_column_counts
+                on a.attnum = constraint_column_counts.attnum
+                and a.attrelid = constraint_column_counts.rel_oid
+            left join
+                fkey_references
+                on a.attnum = fkey_references.attnum
+                and a.attrelid = fkey_references.rel_oid
+            where
+                c.relname = '{item.label}'
+                and n.nspname = '{item.parent.label}'
+                and a.attnum > 0
+                and not a.attisdropped
+            order by a.attnum
+            """.strip("\n")
+        )
+    )
+
+
+def show_describe_table_indexes(
+    item: "RelationCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    if item.parent is None:
+        driver.notify(
+            f"Could not describe {item.label} due to missing schema reference.",
+            severity="error",
+        )
+        return
+    driver.insert_text_in_new_buffer(
+        dedent(
+            f"""
+            with
+                index_columns as (
+                    select
+                        i.indexrelid, c.oid as rel_oid, unnest(i.indkey) as attnum
+                    from pg_catalog.pg_index i
+                    join pg_catalog.pg_class c on c.oid = i.indrelid
+                    where c.relname = '{item.label}'
+                ),
+                index_column_names as (
+                    select
+                        index_columns.indexrelid, 
+                        string_agg(pg_attribute.attname, ', ') as columns
+                    from index_columns
+                    join
+                        pg_catalog.pg_attribute
+                        on index_columns.rel_oid = pg_attribute.attrelid
+                        and index_columns.attnum = pg_attribute.attnum
+                    group by 1
+                )
+
+            select
+                c.relname as "Table",
+                i.indexrelid::regclass::text as "Index Name",
+                pg_am.amname as "Index Type",
+                index_column_names.columns as "Columns",
+                i.indisprimary as "Is PK",
+                i.indisunique as "Is Unique",
+                i.indisclustered as "Is Clustered",
+                i.indisvalid as "Is Valid"
+            from pg_catalog.pg_index i
+            join pg_catalog.pg_class c on c.oid = i.indrelid
+            left join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+            join pg_catalog.pg_class ic on i.indexrelid = ic.oid
+            join pg_catalog.pg_am on ic.relam = pg_am.oid
+            left join
+                index_column_names on i.indexrelid = index_column_names.indexrelid
+            where
+                c.relname = '{item.label}'
+                and n.nspname = '{item.parent.label}'
+            """.strip("\n")
+        )
+    )
+
+
+def show_describe_table_constraints(
+    item: "RelationCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    if item.parent is None:
+        driver.notify(
+            f"Could not describe {item.label} due to missing schema reference.",
+            severity="error",
+        )
+        return
+    driver.insert_text_in_new_buffer(
+        dedent(
+            f"""
+            with
+                constraint_columns as (
+                    select con.oid, c.oid as rel_oid, unnest(con.conkey) as attnum
+                    from pg_catalog.pg_constraint con
+                    join pg_catalog.pg_class c on con.conrelid = c.oid
+                    where c.relname = '{item.label}'
+                ),
+                constraint_column_names as (
+                    select
+                        constraint_columns.oid,
+                        string_agg(pg_attribute.attname, ', ') as columns
+                    from constraint_columns
+                    join
+                        pg_catalog.pg_attribute
+                        on constraint_columns.rel_oid = pg_attribute.attrelid
+                        and constraint_columns.attnum = pg_attribute.attnum
+                    group by 1
+                ),
+                constraint_foreign_columns as (
+                    select con.oid, c.oid as rel_oid, unnest(con.confkey) as attnum
+                    from pg_catalog.pg_constraint con
+                    join pg_catalog.pg_class c on con.conrelid = c.oid
+                    where c.relname = '{item.label}'
+                ),
+                constraint_foreign_column_names as (
+                    select
+                        constraint_foreign_columns.oid,
+                        string_agg(pg_attribute.attname, ', ') as columns
+                    from constraint_foreign_columns
+                    join
+                        pg_catalog.pg_attribute
+                        on constraint_foreign_columns.rel_oid = pg_attribute.attrelid
+                        and constraint_foreign_columns.attnum = pg_attribute.attnum
+                    group by 1
+                )
+
+            select
+                c.relname as "Table",
+                con.conname as "Constraint name",
+                case
+                    con.contype
+                    when 'c'
+                    then 'Check'
+                    when 'n'
+                    then 'Not Null'
+                    when 'p'
+                    then 'Primary Key'
+                    when 'f'
+                    then 'Foreign Key'
+                    when 'u'
+                    then 'Unique'
+                    when 't'
+                    then 'Trigger'
+                    when 'x'
+                    then 'Exclusion'
+                    else con.contype::text
+                end as "Constraint Type",
+                constraint_column_names.columns as "Columns",
+                c.relnamespace::regnamespace
+                || '.'
+                || fc.relname
+                || '('
+                || constraint_foreign_column_names.columns
+                || ')' as "References",
+                con.conislocal as "Is Local",
+                con.convalidated as "Is Validated",
+                case
+                    con.confupdtype
+                    when 'a'
+                    then 'no action'
+                    when 'r'
+                    then 'restrict'
+                    when 'c'
+                    then 'cascade'
+                    when 'n'
+                    then 'set null'
+                    when 'd'
+                    then 'set default'
+                    else con.confupdtype::text
+                end as "FK Update Type",
+                case
+                    con.confdeltype
+                    when 'a'
+                    then 'no action'
+                    when 'r'
+                    then 'restrict'
+                    when 'c'
+                    then 'cascade'
+                    when 'n'
+                    then 'set null'
+                    when 'd'
+                    then 'set default'
+                    else con.confdeltype::text
+                end as "FK Delete Type",
+                case
+                    con.confmatchtype
+                    when 's'
+                    then 'simple'
+                    when 'f'
+                    then 'full'
+                    when 'p'
+                    then 'partial'
+                    else con.confmatchtype::text
+                end as "FK Match Type"
+            from pg_catalog.pg_constraint con
+            join pg_catalog.pg_class c on con.conrelid = c.oid
+            join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+            left join pg_catalog.pg_class fc on con.confrelid = fc.oid
+            left join
+                constraint_column_names on con.oid = constraint_column_names.oid
+            left join
+                constraint_foreign_column_names
+                on con.oid = constraint_foreign_column_names.oid
+            where
+                c.relname = '{item.label}'
+                and n.nspname = '{item.parent.label}'
+            """.strip("\n")
+        )
+    )
+
+
+def show_view_definition(
+    item: "ViewCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    if item.connection is None or item.parent is None:
+        return
+    view_def_query = f"""
+        select pg_catalog.pg_get_viewdef(c.oid, true)
+        from pg_catalog.pg_class as c
+        left join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+        where c.relname = '{item.label}' and n.nspname = '{item.parent.label}'
+        """.strip("\n")
+    cur = item.connection.execute(view_def_query)
+    if cur is None:
+        return
+    result = cur.fetchall()
+    if result is None:
+        return
+    view_def: str = result[0][0]
+    driver.insert_text_in_new_buffer(
+        f"-- View definition for {item.query_name}\n" + view_def
+    )
+
+
 def insert_columns_at_cursor(
     item: "RelationCatalogItem",
     driver: "HarlequinDriver",