fix(plugin-postgresql): list a partitioned table once and nest its partitions (#1984)

CHANGELOG.md

@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- A PostgreSQL partitioned table is listed once instead of once per partition. Expand it in the sidebar to see its partitions, and expand one again if it is subpartitioned. A table split into hundreds of partitions no longer buries the rest of the schema, and those partitions no longer fill up query autocomplete, Cmd+K, and the export picker. (#1984)
 - An SSH tunnel that cannot reach the database now says so, instead of leaving the database driver to report a timeout reading the server greeting. MySQL showed this as "Lost connection to server at 'handshake: reading initial communication packet', system error: 35", which named no cause. TablePro now checks the destination is reachable before the tunnel is handed over, and reports whether it was refused or timed out. The most common cause is the **Host** field holding the server's public address while the database only listens on `127.0.0.1`; the SSH server resolves that address from where it sits, so **Host** should be `localhost`. (#1981)
 - The macOS local network prompt now says why TablePro wants access. The app never declared a reason, so the prompt showed a generic message that was easy to deny, which left SSH tunnels and databases on your own network failing with "no route to host".
 - An SSH tunnel no longer dies while it is busy. The keep-alive read a send that had to wait as a dead tunnel, so heavy traffic through the tunnel could tear it down along with every connection using it.

Plugins/PostgreSQLDriverPlugin/PostgreSQLCapabilities.swift

@@ -17,6 +17,7 @@ struct PostgreSQLCapabilities: Sendable, Equatable {
 
     var hasIdentityColumns: Bool { serverVersion >= 100_000 }
     var hasGeneratedColumns: Bool { serverVersion >= 120_000 }
+    var hasDeclarativePartitioning: Bool { serverVersion >= 100_000 }
 
     var hasArrayPosition: Bool { serverVersion >= 90_500 }
     var hasOrderedAggregates: Bool { serverVersion >= 90_000 }

Plugins/PostgreSQLDriverPlugin/PostgreSQLPluginDriver.swift

@@ -15,8 +15,8 @@ class PostgreSQLPluginDriver: LibPQBackedDriver, @unchecked Sendable {
 
     private static let logger = Logger(subsystem: "com.TablePro.PostgreSQLDriver", category: "PostgreSQLPluginDriver")
 
-    private static let undefinedTableSQLState = "42P01"
-    private static let undefinedFunctionSQLState = "42883"
+    private static let undefinedTableSQLState = PostgreSQLTableListingLadder.undefinedTableSQLState
+    private static let undefinedFunctionSQLState = PostgreSQLTableListingLadder.undefinedFunctionSQLState
 
     private var catalogPresence: PostgreSQLCatalogPresence?
 
@@ -156,29 +156,35 @@ class PostgreSQLPluginDriver: LibPQBackedDriver, @unchecked Sendable {
 
     func fetchTables(schema: String?) async throws -> [PluginTableInfo] {
         let schemaLiteral = escapeLiteral(schema ?? core.currentSchema)
-        func query(includeOptionalCatalogs: Bool, includeComments: Bool) -> String {
+        func query(_ attempt: PostgreSQLTableListingAttempt) -> String {
             PostgreSQLSchemaQueries.fetchTables(
                 schemaLiteral: schemaLiteral,
-                includeMaterializedViews: includeOptionalCatalogs && includesMaterializedViews(),
-                includeForeignTables: includeOptionalCatalogs && includesForeignTables(),
-                includeComments: includeComments
+                includeMaterializedViews: attempt.includeOptionalCatalogs && includesMaterializedViews(),
+                includeForeignTables: attempt.includeOptionalCatalogs && includesForeignTables(),
+                includeComments: attempt.includeComments,
+                includePartitionAwareness: attempt.includePartitionAwareness
             )
         }
 
-        let result: PluginQueryResult
-        do {
-            result = try await execute(query: query(includeOptionalCatalogs: true, includeComments: true))
-        } catch let error as LibPQPluginError where error.sqlState == Self.undefinedTableSQLState {
-            result = try await execute(query: query(includeOptionalCatalogs: false, includeComments: true))
-        } catch let error as LibPQPluginError where error.sqlState == Self.undefinedFunctionSQLState {
-            result = try await execute(query: query(includeOptionalCatalogs: false, includeComments: false))
+        var result: PluginQueryResult?
+        for attempt in PostgreSQLTableListingLadder.degradableAttempts where result == nil {
+            do {
+                result = try await execute(query: query(attempt))
+            } catch let error as LibPQPluginError where PostgreSQLTableListingLadder.isDegradable(sqlState: error.sqlState) {
+                Self.logger.debug("Table listing degrading past \(attempt.label, privacy: .public): \(error.localizedDescription)")
+            }
+        }
+        if result == nil {
+            result = try await execute(query: query(PostgreSQLTableListingLadder.leastCapableAttempt))
         }
 
+        guard let result else { return [] }
         return result.rows.compactMap { row -> PluginTableInfo? in
             guard let name = row[0].asText else { return nil }
             let typeStr = row[1].asText ?? "BASE TABLE"
             let type: String
             switch typeStr {
+            case "PARTITIONED TABLE": type = "PARTITIONED TABLE"
             case "MATERIALIZED VIEW": type = "MATERIALIZED VIEW"
             case "FOREIGN TABLE":     type = "FOREIGN TABLE"
             case "VIEW":              type = "VIEW"
@@ -189,6 +195,26 @@ class PostgreSQLPluginDriver: LibPQBackedDriver, @unchecked Sendable {
         }
     }
 
+    func fetchPartitions(table: String, schema: String?) async throws -> [PluginTableInfo] {
+        guard versionedCapabilities.hasDeclarativePartitioning else { return [] }
+        let schemaLiteral = escapeLiteral(schema ?? core.currentSchema)
+        let result = try await execute(
+            query: PostgreSQLSchemaQueries.fetchPartitions(
+                schemaLiteral: schemaLiteral,
+                tableLiteral: escapeLiteral(table)
+            )
+        )
+        return result.rows.compactMap { row -> PluginTableInfo? in
+            guard let name = row[0].asText else { return nil }
+            let isSubpartitioned = row[safe: 1]?.asText == "p"
+            return PluginTableInfo(
+                name: name,
+                type: isSubpartitioned ? "PARTITIONED TABLE" : "TABLE",
+                schema: schema ?? core.currentSchema,
+                comment: nil
+            )
+        }
+    }
 
     func fetchIndexes(table: String, schema: String?) async throws -> [PluginIndexInfo] {
         let columnOrdering = versionedCapabilities.hasArrayPosition
@@ -656,10 +682,19 @@ class PostgreSQLPluginDriver: LibPQBackedDriver, @unchecked Sendable {
         let query = """
             SELECT
                 (SELECT COUNT(*)
-                 FROM information_schema.tables
-                 WHERE table_catalog = '\(escapedDbLiteral)'
-                   AND table_schema NOT LIKE 'pg!_%' ESCAPE '!'
-                   AND table_schema <> 'information_schema'),
+                 FROM information_schema.tables t
+                 WHERE t.table_catalog = '\(escapedDbLiteral)'
+                   AND t.table_schema NOT LIKE 'pg!_%' ESCAPE '!'
+                   AND t.table_schema <> 'information_schema'
+                   AND NOT EXISTS (
+                         SELECT 1
+                         FROM pg_catalog.pg_inherits i
+                         JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent
+                         JOIN pg_catalog.pg_class child ON child.oid = i.inhrelid
+                         JOIN pg_catalog.pg_namespace cn ON cn.oid = child.relnamespace
+                         WHERE cn.nspname = t.table_schema
+                           AND child.relname = t.table_name
+                           AND parent.relkind IN ('p', 'I'))),
                 pg_database_size('\(escapedDbLiteral)')
         """
         let result = try await execute(query: query)

Plugins/PostgreSQLDriverPlugin/PostgreSQLSchemaQueries.swift

@@ -76,23 +76,62 @@ enum PostgreSQLSchemaQueries {
     /// `includeComments` projects each table's comment via `obj_description` /
     /// `to_regclass`. Engines that lack those functions fail the whole listing,
     /// so the caller passes `false` to fall back to a comment-free listing.
+    ///
+    /// `includePartitionAwareness` labels a declarative partition parent as
+    /// `PARTITIONED TABLE` and drops its partition children, which
+    /// `information_schema.tables` reports as plain `BASE TABLE` rows
+    /// indistinguishable from the parent. The test is `pg_inherits` joined to
+    /// the parent's `relkind`, not `pg_class.relispartition`: `relispartition`
+    /// only exists from PostgreSQL 10, and referencing a missing column fails
+    /// at parse time, which would break the listing outright on older servers.
+    /// Comparing `relkind` against `'p'`/`'I'` is a value test on a column
+    /// present since PostgreSQL 8, so it parses everywhere and simply matches
+    /// nothing before declarative partitioning existed. Rows still come from
+    /// `information_schema.tables`, which keeps its privilege filtering; the
+    /// catalog joins only label and exclude rows it already returned. The
+    /// caller passes `false` for engines without these catalogs.
+    ///
+    /// Legacy `INHERITS` children stay listed on purpose. Their parent is an
+    /// ordinary table (`relkind = 'r'`), and they are independently useful
+    /// tables rather than an implementation detail of one parent.
     static func fetchTables(
         schemaLiteral: String,
         includeMaterializedViews: Bool,
         includeForeignTables: Bool,
-        includeComments: Bool = true
+        includeComments: Bool = true,
+        includePartitionAwareness: Bool = true
     ) -> String {
         func commentColumn(_ expression: String) -> String {
             includeComments ? expression : "NULL::text"
         }
 
+        let partitionJoin = includePartitionAwareness ? """
+
+            LEFT JOIN pg_catalog.pg_namespace pn ON pn.nspname = t.table_schema
+            LEFT JOIN pg_catalog.pg_class pc ON pc.relnamespace = pn.oid AND pc.relname = t.table_name
+            """ : ""
+
+        let tableTypeColumn = includePartitionAwareness
+            ? "CASE WHEN pc.relkind = 'p' THEN 'PARTITIONED TABLE' ELSE t.table_type END"
+            : "t.table_type"
+
+        let partitionFilter = includePartitionAwareness ? """
+
+              AND NOT EXISTS (
+                    SELECT 1
+                    FROM pg_catalog.pg_inherits i
+                    JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent
+                    WHERE i.inhrelid = pc.oid
+                      AND parent.relkind IN ('p', 'I'))
+            """ : ""
+
         var unions: [String] = [
             """
-            SELECT t.table_name, t.table_type,
+            SELECT t.table_name, \(tableTypeColumn) AS table_type,
                    \(commentColumn("obj_description(to_regclass(quote_ident(t.table_schema) || '.' || quote_ident(t.table_name)), 'pg_class')")) AS table_comment
-            FROM information_schema.tables t
+            FROM information_schema.tables t\(partitionJoin)
             WHERE t.table_schema = '\(schemaLiteral)'
-              AND t.table_type IN ('BASE TABLE', 'VIEW')
+              AND t.table_type IN ('BASE TABLE', 'VIEW')\(partitionFilter)
             """
         ]
 
@@ -123,6 +162,27 @@ enum PostgreSQLSchemaQueries {
         return unions.joined(separator: "\nUNION ALL\n") + "\nORDER BY table_name"
     }
 
+    /// Lists one partitioned table's direct partitions, ordered so the DEFAULT
+    /// partition sorts last. A child that is itself subpartitioned comes back
+    /// with `relkind = 'p'` so it can be expanded in turn.
+    ///
+    /// `relpartbound` exists only from PostgreSQL 10, so unlike `fetchTables`
+    /// this query cannot be issued against an older server. The caller gates it
+    /// on `PostgreSQLCapabilities.hasDeclarativePartitioning`.
+    static func fetchPartitions(schemaLiteral: String, tableLiteral: String) -> String {
+        """
+        SELECT cc.relname, cc.relkind
+        FROM pg_catalog.pg_inherits i
+        JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent
+        JOIN pg_catalog.pg_namespace pn ON pn.oid = parent.relnamespace
+        JOIN pg_catalog.pg_class cc ON cc.oid = i.inhrelid
+        WHERE pn.nspname = '\(schemaLiteral)'
+          AND parent.relname = '\(tableLiteral)'
+          AND parent.relkind = 'p'
+        ORDER BY pg_catalog.pg_get_expr(cc.relpartbound, cc.oid) = 'DEFAULT', cc.relname
+        """
+    }
+
     static func setSearchPath(toSchema schema: String) -> String {
         let quotedIdentifier = "\"\(schema.replacingOccurrences(of: "\"", with: "\"\""))\""
         return "SET search_path TO \(quotedIdentifier)"

Plugins/PostgreSQLDriverPlugin/PostgreSQLTableListingLadder.swift

@@ -0,0 +1,51 @@
+//
+//  PostgreSQLTableListingLadder.swift
+//  PostgreSQLDriverPlugin
+//
+//  Ordered fallbacks for the table listing, extracted so the ordering can be
+//  exercised by unit tests via TableProTests/PluginTestSources.
+//
+
+import Foundation
+
+struct PostgreSQLTableListingAttempt: Sendable, Equatable {
+    let includeOptionalCatalogs: Bool
+    let includeComments: Bool
+    let includePartitionAwareness: Bool
+
+    var label: String {
+        "catalogs=\(includeOptionalCatalogs) comments=\(includeComments) partitions=\(includePartitionAwareness)"
+    }
+}
+
+/// A PostgreSQL-compatible engine may lack the optional catalogs the listing
+/// reads, so each rung drops one more of them and retries. Partition awareness
+/// degrades last: it is the only rung that changes which rows come back rather
+/// than which columns, so a partition-free listing is worth every other
+/// fallback first. A single flat `catch` cannot express this, because the same
+/// SQLSTATE means a different missing catalog at each rung.
+enum PostgreSQLTableListingLadder {
+    static let undefinedTableSQLState = "42P01"
+    static let undefinedFunctionSQLState = "42883"
+
+    static let degradableAttempts: [PostgreSQLTableListingAttempt] = [
+        PostgreSQLTableListingAttempt(
+            includeOptionalCatalogs: true, includeComments: true, includePartitionAwareness: true
+        ),
+        PostgreSQLTableListingAttempt(
+            includeOptionalCatalogs: false, includeComments: true, includePartitionAwareness: true
+        ),
+        PostgreSQLTableListingAttempt(
+            includeOptionalCatalogs: false, includeComments: false, includePartitionAwareness: true
+        )
+    ]
+
+    static let leastCapableAttempt = PostgreSQLTableListingAttempt(
+        includeOptionalCatalogs: false, includeComments: false, includePartitionAwareness: false
+    )
+
+    static func isDegradable(sqlState: String?) -> Bool {
+        guard let sqlState else { return false }
+        return sqlState == undefinedTableSQLState || sqlState == undefinedFunctionSQLState
+    }
+}

Plugins/TableProPluginKit/PluginDatabaseDriver.swift

@@ -82,6 +82,7 @@ public protocol PluginDatabaseDriver: AnyObject, Sendable {
     func executeUserQuery(query: String, rowCap: Int?, parameters: [PluginCellValue]?) async throws -> PluginQueryResult
 
     func fetchTables(schema: String?) async throws -> [PluginTableInfo]
+    func fetchPartitions(table: String, schema: String?) async throws -> [PluginTableInfo]
     func fetchColumns(table: String, schema: String?) async throws -> [PluginColumnInfo]
     func fetchIndexes(table: String, schema: String?) async throws -> [PluginIndexInfo]
     func fetchForeignKeys(table: String, schema: String?) async throws -> [PluginForeignKeyInfo]
@@ -201,6 +202,10 @@ public extension PluginDatabaseDriver {
 
     func fetchTriggers(table: String, schema: String?) async throws -> [PluginTriggerInfo] { [] }
 
+    /// Engines whose partitions are metadata on one table object, rather than
+    /// separate relations, have nothing to nest and keep the empty default.
+    func fetchPartitions(table: String, schema: String?) async throws -> [PluginTableInfo] { [] }
+
     func createTriggerTemplate(table: String, schema: String?) -> String? { nil }
     func fetchTriggerDefinition(name: String, table: String, schema: String?) async throws -> String? { nil }
     func generateDropTriggerSQL(name: String, table: String, schema: String?) -> String? { nil }

TablePro/Core/Database/DatabaseDriver.swift

@@ -69,6 +69,9 @@ protocol DatabaseDriver: AnyObject, Sendable {
 
     func fetchTables(schema: String?) async throws -> [TableInfo]
 
+    /// Fetch the direct partitions of one partitioned table
+    func fetchPartitions(table: String, schema: String?) async throws -> [TableInfo]
+
     /// Fetch columns for a specific table
     func fetchColumns(table: String) async throws -> [ColumnInfo]
 
@@ -252,6 +255,8 @@ extension DatabaseDriver {
         try await fetchColumns(table: table)
     }
 
+    func fetchPartitions(table: String, schema: String?) async throws -> [TableInfo] { [] }
+
     func fetchTriggers(table: String) async throws -> [TriggerInfo] { [] }
 
     func createTriggerTemplate(table: String) -> String? { nil }

TablePro/Core/Database/TableOperationSQLBuilder.swift

@@ -111,7 +111,7 @@ struct TableOperationSQLBuilder {
             return "MATERIALIZED VIEW"
         case .foreignTable:
             return "FOREIGN TABLE"
-        case .table, .systemTable, .none:
+        case .table, .systemTable, .partitionedTable, .none:
             return "TABLE"
         }
     }

TablePro/Core/Plugins/PluginDriverAdapter.swift

@@ -181,11 +181,19 @@ final class PluginDriverAdapter: DatabaseDriver, SchemaSwitchable {
         return pluginTables.map { mapPluginTable($0, schemaFallback: resolvedSchema) }
     }
 
+    func fetchPartitions(table: String, schema: String?) async throws -> [TableInfo] {
+        let resolvedSchema = schema ?? pluginDriver.currentSchema
+        let pluginTables = try await pluginDriver.fetchPartitions(table: table, schema: resolvedSchema)
+        return pluginTables.map { mapPluginTable($0, schemaFallback: resolvedSchema) }
+    }
+
     private func mapPluginTable(_ table: PluginTableInfo, schemaFallback: String?) -> TableInfo {
         let tableType: TableInfo.TableType
         switch table.type.lowercased() {
         case "table", "base table", "prefix":
             tableType = .table
+        case "partitioned table", "partitioned_table":
+            tableType = .partitionedTable
         case "view":
             tableType = .view
         case "materialized view", "materialized_view":