Option to exclude partitions (PostgreSQL)

[amacfie]

It would be nice to be able to exclude partitions but include partitioned tables. Reference: https://www.postgresql.org/docs/current/ddl-partitioning.html

[amacfie]

Based on https://dba.stackexchange.com/questions/40441/get-all-partition-names-for-a-table, it looks like the following lists all tables that are a part of a partitioned table:

SELECT
    nmsp_child.nspname  AS child_schema,
    child.relname       AS child
FROM pg_inherits
    JOIN pg_class parent            ON pg_inherits.inhparent = parent.oid
    JOIN pg_class child             ON pg_inherits.inhrelid   = child.oid
    JOIN pg_namespace nmsp_parent   ON nmsp_parent.oid  = parent.relnamespace
    JOIN pg_namespace nmsp_child    ON nmsp_child.oid   = child.relnamespace
WHERE parent.relkind = 'p';

[amacfie]

I'm not sure if this should be built into sqlacodegen so here's a script that generates a value for the tables parameter of sqlacodegen that excludes partitions:

#!/usr/bin/env python3
"""r
Prints a comma-separated list of all tables (and optionally views) that aren't
partitions of a partitioned table.
"""

# /// script
# dependencies = [
#   "typer==0.9.*",
#   "asyncpg==0.29.*",
# ]
# requires-python = ">=3.9"
# ///

import asyncio
from typing import Annotated

import asyncpg
import typer


partition_q = r"""
SELECT
    child.relname
FROM pg_inherits
    JOIN pg_class parent            ON pg_inherits.inhparent = parent.oid
    JOIN pg_class child             ON pg_inherits.inhrelid   = child.oid
    JOIN pg_namespace nmsp_child    ON nmsp_child.oid   = child.relnamespace
WHERE parent.relkind = 'p' AND nmsp_child.nspname = $1;
"""

table_q = r"""
SELECT table_name FROM information_schema.tables 
WHERE
    table_schema = $1
    AND ((table_type = 'BASE TABLE') OR $2);
"""


async def get_tables(dsn: str, schema: str, include_views: bool):
    conn = await asyncpg.connect(dsn)
    tables = {
        rec[0] for rec in await conn.fetch(table_q, schema, include_views)
    }
    parts = {rec[0] for rec in await conn.fetch(partition_q, schema)}
    return ",".join(sorted(tables - parts))


def main(
    dsn: Annotated[
        str, typer.Argument(help="postgres://user:password@host:port/database")
    ],
    schema: str = "public",
    include_views: bool = True,
):
    # the "tables" param in sqlacodegen applies to both tables and views
    # so we have to optionally include views.

    tables = asyncio.run(get_tables(dsn, schema, include_views))
    print(tables)


if __name__ == "__main__":
    typer.run(main)