598 lines
21 KiB
Python
598 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
from textwrap import dedent
|
|
from typing import TYPE_CHECKING, Literal, Sequence
|
|
|
|
from harlequin.catalog import CatalogItem
|
|
from harlequin.exception import HarlequinQueryError
|
|
|
|
if TYPE_CHECKING:
|
|
from harlequin.driver import HarlequinDriver
|
|
|
|
from harlequin_postgres.catalog import (
|
|
ColumnCatalogItem,
|
|
DatabaseCatalogItem,
|
|
RelationCatalogItem,
|
|
SchemaCatalogItem,
|
|
ViewCatalogItem,
|
|
)
|
|
|
|
|
|
def execute_use_statement(
|
|
item: "SchemaCatalogItem",
|
|
driver: "HarlequinDriver",
|
|
) -> None:
|
|
if item.connection is None:
|
|
return
|
|
try:
|
|
item.connection.execute(f"set search_path to {item.qualified_identifier}")
|
|
except HarlequinQueryError:
|
|
driver.notify("Could not switch context", severity="error")
|
|
raise
|
|
else:
|
|
driver.notify(f"Editor context switched to {item.label}")
|
|
|
|
|
|
def execute_drop_schema_statement(
|
|
item: "SchemaCatalogItem",
|
|
driver: "HarlequinDriver",
|
|
) -> None:
|
|
def _drop_schema() -> None:
|
|
if item.connection is None:
|
|
return
|
|
try:
|
|
item.connection.execute(f"drop schema {item.qualified_identifier} cascade")
|
|
except HarlequinQueryError:
|
|
driver.notify(f"Could not drop schema {item.label}", severity="error")
|
|
raise
|
|
else:
|
|
driver.notify(f"Dropped schema {item.label}")
|
|
driver.refresh_catalog()
|
|
|
|
if item.children or item.fetch_children():
|
|
driver.confirm_and_execute(callback=_drop_schema)
|
|
else:
|
|
_drop_schema()
|
|
|
|
|
|
def execute_drop_database_statement(
|
|
item: "DatabaseCatalogItem",
|
|
driver: "HarlequinDriver",
|
|
) -> None:
|
|
def _drop_database() -> None:
|
|
if item.connection is None:
|
|
return
|
|
try:
|
|
item.connection.execute(f"drop database {item.qualified_identifier}")
|
|
except HarlequinQueryError:
|
|
driver.notify(f"Could not drop database {item.label}", severity="error")
|
|
raise
|
|
else:
|
|
driver.notify(f"Dropped database {item.label}")
|
|
driver.refresh_catalog()
|
|
|
|
if item.children or item.fetch_children():
|
|
driver.confirm_and_execute(callback=_drop_database)
|
|
else:
|
|
_drop_database()
|
|
|
|
|
|
def execute_drop_relation_statement(
|
|
item: "RelationCatalogItem",
|
|
driver: "HarlequinDriver",
|
|
relation_type: Literal["view", "table", "foreign table"],
|
|
) -> None:
|
|
def _drop_relation() -> None:
|
|
if item.connection is None:
|
|
return
|
|
try:
|
|
item.connection.execute(f"drop {relation_type} {item.qualified_identifier}")
|
|
except HarlequinQueryError:
|
|
driver.notify(
|
|
f"Could not drop {relation_type} {item.label}", severity="error"
|
|
)
|
|
raise
|
|
else:
|
|
driver.notify(f"Dropped {relation_type} {item.label}")
|
|
driver.refresh_catalog()
|
|
|
|
driver.confirm_and_execute(callback=_drop_relation)
|
|
|
|
|
|
def execute_drop_table_statement(
|
|
item: "RelationCatalogItem", driver: "HarlequinDriver"
|
|
) -> None:
|
|
execute_drop_relation_statement(item=item, driver=driver, relation_type="table")
|
|
|
|
|
|
def execute_drop_foreign_table_statement(
|
|
item: "RelationCatalogItem", driver: "HarlequinDriver"
|
|
) -> None:
|
|
execute_drop_relation_statement(
|
|
item=item, driver=driver, relation_type="foreign table"
|
|
)
|
|
|
|
|
|
def execute_drop_view_statement(
|
|
item: "RelationCatalogItem", driver: "HarlequinDriver"
|
|
) -> None:
|
|
execute_drop_relation_statement(item=item, driver=driver, relation_type="view")
|
|
|
|
|
|
def show_select_star(
|
|
item: "RelationCatalogItem",
|
|
driver: "HarlequinDriver",
|
|
) -> None:
|
|
driver.insert_text_in_new_buffer(
|
|
dedent(
|
|
f"""
|
|
select *
|
|
from {item.qualified_identifier}
|
|
limit 100
|
|
""".strip("\n")
|
|
)
|
|
)
|
|
|
|
|
|
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",
|
|
) -> None:
|
|
if item.loaded:
|
|
cols: Sequence["CatalogItem" | "ColumnCatalogItem"] = item.children
|
|
else:
|
|
cols = item.fetch_children()
|
|
driver.insert_text_at_selection(text=",\n".join(c.query_name for c in cols))
|