fix(plugin-postgresql): resolve the oid of a type created after connect from the catalog

CHANGELOG.md

@@ -27,6 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Last line of a helper process's output lost when it exits right after writing it.
 - Composite, range and extension-typed PostgreSQL columns labelled `ENUM(…)` in the structure editor.
+- Columns of a PostgreSQL enum created during the session shown as text until the next reconnect.
 - Sidebar routines, triggers and types from the previous database after switching while it was still loading.
 - Silent fallback order when foreign keys between the exported tables form a cycle. (#2517)
 - Foreign keys declared twice in a SQL export on MySQL, SQL Server, DuckDB, Snowflake, CockroachDB and Redshift, and as an unsupported `ALTER TABLE` on SQLite, libSQL and Cloudflare D1. (#2517)

Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift

@@ -177,8 +177,8 @@ final class LibPQDriverCore: @unchecked Sendable {
         libpqConnection?.setPostgisOidMap(map)
     }
 
-    func setEnumOidMap(_ map: [UInt32: String]) {
-        libpqConnection?.setEnumOidMap(map)
+    func mergeCatalogTypeNames(_ names: [UInt32: String]) {
+        libpqConnection?.mergeCatalogTypeNames(names)
     }
 
     func applyQueryTimeout(_ seconds: Int) async throws {

Plugins/PostgreSQLDriverPlugin/LibPQPluginConnection.swift

@@ -47,7 +47,7 @@ struct LibPQPluginQueryResult {
 
 // MARK: - Type Mapping
 
-private func pgOidToTypeName(_ oid: UInt32) -> String {
+private func pgOidToTypeName(_ oid: UInt32) -> String? {
     switch oid {
     case 16: return "boolean"
     case 17: return "bytea"
@@ -103,7 +103,7 @@ private func pgOidToTypeName(_ oid: UInt32) -> String {
     case 1_700: return "numeric"
     case 2_950: return "uuid"
     case 3_802: return "jsonb"
-    default: return "unknown"
+    default: return nil
     }
 }
 
@@ -133,7 +133,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
     private var _cachedServerVersionNumber: Int32 = 0
     private var _isConnectCancelled: Bool = false
     private var _postgisOidMap: [UInt32: String] = [:]
-    private var _enumOidMap: [UInt32: String] = [:]
+    private var _catalogTypeNames: [UInt32: String] = [:]
 
     var isConnected: Bool {
         stateLock.lock()
@@ -403,17 +403,89 @@ final class LibPQPluginConnection: @unchecked Sendable {
         return _postgisOidMap
     }
 
-    func setEnumOidMap(_ map: [UInt32: String]) {
+    func mergeCatalogTypeNames(_ names: [UInt32: String]) {
         stateLock.lock()
-        _enumOidMap = map
+        _catalogTypeNames.merge(names) { _, learned in learned }
         stateLock.unlock()
     }
 
     private func resolveTypeName(_ oid: UInt32) -> String {
         stateLock.lock()
-        let mapped = _enumOidMap[oid]
+        let mapped = _catalogTypeNames[oid]
         stateLock.unlock()
-        return mapped ?? pgOidToTypeName(oid)
+        return mapped ?? pgOidToTypeName(oid) ?? PostgreSQLCatalogTypeNames.unresolved
+    }
+
+    private func unresolvedOids(in oids: [UInt32]) -> [UInt32] {
+        stateLock.lock()
+        defer { stateLock.unlock() }
+        return oids.filter { _catalogTypeNames[$0] == nil && pgOidToTypeName($0) == nil }
+    }
+
+    /// A type created after connect has an oid the connect-time probe never saw, so its columns
+    /// came back as text until the next reconnect. The oids a result leaves unresolved are looked
+    /// up on the same connection once the result is fully read, which is the only moment libpq
+    /// allows another statement, and remembered for every later result. An oid the catalog does
+    /// not know is remembered as unresolved for the same reason. A lookup that fails inside an
+    /// aborted transaction remembers nothing, because it will succeed after the rollback.
+    private func learnTypeNames(for oids: [UInt32], conn: OpaquePointer) {
+        guard let query = PostgreSQLCatalogTypeNames.lookupQuery(oids: oids) else { return }
+        let result: OpaquePointer? = query.withCString { PQexec(conn, $0) }
+        guard let result else { return }
+        defer { PQclear(result) }
+
+        guard PQresultStatus(result) == PGRES_TUPLES_OK else {
+            guard getResultError(from: result).sqlState != Self.transactionAbortedSQLState else { return }
+            mergeCatalogTypeNames(PostgreSQLCatalogTypeNames.names(for: oids, rows: []))
+            return
+        }
+        mergeCatalogTypeNames(PostgreSQLCatalogTypeNames.names(for: oids, rows: Self.textRows(from: result)))
+    }
+
+    /// A streaming result sends its header with the first row, before anything could be looked
+    /// up, so a `SELECT` run right after a `CREATE TYPE` in the same tab would still read the new
+    /// enum as text once. The statement's own command tag says a type was just created, and one
+    /// enum probe there puts the oid in place before the next statement is sent.
+    private func noteCommandTag(_ tag: String?, conn: OpaquePointer) {
+        guard tag == Self.createTypeCommandTag else { return }
+        let query = PostgreSQLSchemaQueries.enumTypeOidQuery
+        let result: OpaquePointer? = query.withCString { PQexec(conn, $0) }
+        guard let result else { return }
+        defer { PQclear(result) }
+        guard PQresultStatus(result) == PGRES_TUPLES_OK else { return }
+        mergeCatalogTypeNames(PostgreSQLCatalogTypeNames.enumProbeNames(rows: Self.textRows(from: result)))
+    }
+
+    private static func textRows(from result: OpaquePointer) -> [[String?]] {
+        let numRows = Int(PQntuples(result))
+        let numFields = Int(PQnfields(result))
+        var rows: [[String?]] = []
+        rows.reserveCapacity(numRows)
+        for rowIndex in 0..<numRows {
+            rows.append((0..<numFields).map { fieldIndex in
+                guard PQgetisnull(result, Int32(rowIndex), Int32(fieldIndex)) == 0,
+                      let valuePtr = PQgetvalue(result, Int32(rowIndex), Int32(fieldIndex)) else { return nil }
+                return String(cString: valuePtr)
+            })
+        }
+        return rows
+    }
+
+    private static let transactionAbortedSQLState = "25P02"
+    private static let createTypeCommandTag = "CREATE TYPE"
+
+    private func resolvingUnknownTypes(
+        _ metadata: ColumnMetadata,
+        conn: OpaquePointer
+    ) -> ColumnMetadata {
+        let missing = unresolvedOids(in: metadata.columnOids)
+        guard !missing.isEmpty else { return metadata }
+        learnTypeNames(for: missing, conn: conn)
+        return ColumnMetadata(
+            columns: metadata.columns,
+            columnOids: metadata.columnOids,
+            columnTypeNames: metadata.columnOids.map(resolveTypeName)
+        )
     }
 
     // MARK: - Query Cancellation
@@ -509,6 +581,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
             let affected = getAffectedRows(from: result)
             let cmdTag = getCommandTag(from: result)
             PQclear(result)
+            noteCommandTag(cmdTag, conn: conn)
             return LibPQPluginQueryResult(
                 columns: [],
                 columnOids: [],
@@ -521,7 +594,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
 
         case PGRES_TUPLES_OK:
             defer { PQclear(result) }
-            return try fetchResults(from: result, generation: generation)
+            return try fetchResults(from: result, conn: conn, generation: generation)
 
         default:
             let error = getResultError(from: result)
@@ -646,10 +719,12 @@ final class LibPQPluginConnection: @unchecked Sendable {
 
         if truncated { rows.removeLast() }
 
+        noteCommandTag(commandTag, conn: conn)
+        let resolvedMetadata = metadata.map { resolvingUnknownTypes($0, conn: conn) }
         let bounded = LibPQPluginQueryResult(
-            columns: metadata?.columns ?? [],
-            columnOids: metadata?.columnOids ?? [],
-            columnTypeNames: metadata?.columnTypeNames ?? [],
+            columns: resolvedMetadata?.columns ?? [],
+            columnOids: resolvedMetadata?.columnOids ?? [],
+            columnTypeNames: resolvedMetadata?.columnTypeNames ?? [],
             rows: rows,
             affectedRows: affectedRows,
             commandTag: commandTag,
@@ -744,6 +819,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
             let affected = getAffectedRows(from: result)
             let cmdTag = getCommandTag(from: result)
             PQclear(result)
+            noteCommandTag(cmdTag, conn: conn)
             return LibPQPluginQueryResult(
                 columns: [],
                 columnOids: [],
@@ -756,7 +832,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
 
         case PGRES_TUPLES_OK:
             defer { PQclear(result) }
-            return try fetchResults(from: result, generation: generation)
+            return try fetchResults(from: result, conn: conn, generation: generation)
 
         default:
             let error = getResultError(from: result)
@@ -832,6 +908,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
 
                 var headerSent = false
                 var columnOids: [UInt32] = []
+                var lastCommandTag: String?
                 let batchSize = 5_000
                 var batch: [PluginRow] = []
                 batch.reserveCapacity(batchSize)
@@ -899,6 +976,7 @@ final class LibPQPluginConnection: @unchecked Sendable {
                         PQclear(result)
                         break
                     } else if status == PGRES_COMMAND_OK {
+                        lastCommandTag = getCommandTag(from: result)
                         PQclear(result)
                         break
                     } else {
@@ -919,15 +997,26 @@ final class LibPQPluginConnection: @unchecked Sendable {
                 }
 
                 while let res = PQgetResult(conn) { PQclear(res) }
+                /// The header went out with the first row, so this stream keeps what it said;
+                /// the lookup is for the results that follow.
+                let missing = unresolvedOids(in: columnOids)
+                if !missing.isEmpty {
+                    learnTypeNames(for: missing, conn: conn)
+                }
+                noteCommandTag(lastCommandTag, conn: conn)
                 continuation.finish()
             }
         }
     }
 
     // MARK: - Result Parsing
 
-    private func fetchResults(from result: OpaquePointer, generation: Int) throws -> LibPQPluginQueryResult {
-        let metadata = readColumnMetadata(from: result)
+    private func fetchResults(
+        from result: OpaquePointer,
+        conn: OpaquePointer,
+        generation: Int
+    ) throws -> LibPQPluginQueryResult {
+        let metadata = resolvingUnknownTypes(readColumnMetadata(from: result), conn: conn)
         let parsed = try parseRows(
             from: result,
             columns: metadata.columns,

Plugins/PostgreSQLDriverPlugin/PostgreSQLCatalogTypeNames.swift

@@ -0,0 +1,102 @@
+//
+//  PostgreSQLCatalogTypeNames.swift
+//  PostgreSQLDriver
+//
+
+import Foundation
+
+/// The spelling a result column gets when its type oid is not in libpq's built-in table: an
+/// enum is `ENUM(name)` and an enum array `ENUM[](name)`, which is what the column classifier
+/// reads, a domain is its base type so the cell edits as one, and anything else is its own name.
+/// Both the connect-time enum probe and the per-result lookup spell through here, so a column
+/// reads the same whether its type existed at connect or was created a moment ago.
+enum PostgreSQLCatalogTypeNames {
+    static let unresolved = "unknown"
+
+    static func enumTypeName(_ name: String) -> String { "ENUM(\(name))" }
+
+    static func enumArrayTypeName(_ name: String) -> String { "ENUM[](\(name))" }
+
+    struct Row: Equatable {
+        let oid: UInt32
+        let name: String
+        let kind: Character
+        let domainBase: String?
+        let elementName: String?
+        let elementKind: Character?
+        let elementDomainBase: String?
+    }
+
+    /// An array is `typelem` set on a variable-length type rather than `typcategory = 'A'`,
+    /// because Redshift's catalog predates `typcategory` and the predicate agrees with it on
+    /// every real type: `point` and `name` carry a `typelem` at a fixed length, and the one
+    /// disagreement is the pseudo-type `_record`, which no column has.
+    static func lookupQuery(oids: [UInt32]) -> String? {
+        let list = Set(oids).sorted().map(String.init).joined(separator: ", ")
+        guard !list.isEmpty else { return nil }
+        return """
+            SELECT t.oid::text, t.typname, t.typtype::text,
+                   CASE WHEN t.typtype = 'd' THEN pg_catalog.format_type(t.typbasetype, t.typtypmod) END,
+                   el.typname, el.typtype::text,
+                   CASE WHEN el.typtype = 'd' THEN pg_catalog.format_type(el.typbasetype, el.typtypmod) END
+            FROM pg_catalog.pg_type t
+            LEFT JOIN pg_catalog.pg_type el ON el.oid = t.typelem AND t.typlen = -1
+            WHERE t.oid IN (\(list))
+            """
+    }
+
+    static func row(fromColumns columns: [String?]) -> Row? {
+        guard columns.count >= 7,
+              let oid = columns[0].flatMap({ UInt32($0) }),
+              let name = columns[1],
+              let kind = columns[2]?.first else { return nil }
+        return Row(
+            oid: oid,
+            name: name,
+            kind: kind,
+            domainBase: columns[3],
+            elementName: columns[4],
+            elementKind: columns[5]?.first,
+            elementDomainBase: columns[6]
+        )
+    }
+
+    static func typeName(for row: Row) -> String {
+        if let elementName = row.elementName {
+            guard row.elementKind != "e" else { return enumArrayTypeName(elementName) }
+            return "\(row.elementDomainBase ?? elementName)[]"
+        }
+        if row.kind == "e" { return enumTypeName(row.name) }
+        if row.kind == "d", let base = row.domainBase { return base }
+        return row.name
+    }
+
+    /// The connect-time probe and the refresh after a `CREATE TYPE` both read every enum with its
+    /// scalar and array oids; a zero array oid is a server that does not give enums one.
+    static func enumProbeNames(rows: [[String?]]) -> [UInt32: String] {
+        var names: [UInt32: String] = [:]
+        for row in rows {
+            guard row.count >= 3,
+                  let scalarOid = row[0].flatMap({ UInt32($0) }),
+                  let typeName = row[2] else { continue }
+            names[scalarOid] = enumTypeName(typeName)
+            if let arrayOid = row[1].flatMap({ UInt32($0) }), arrayOid != 0 {
+                names[arrayOid] = enumArrayTypeName(typeName)
+            }
+        }
+        return names
+    }
+
+    /// Every oid asked about gets an entry, so an oid the catalog does not know is not asked
+    /// about again on every following result.
+    static func names(for oids: [UInt32], rows: [[String?]]) -> [UInt32: String] {
+        var names: [UInt32: String] = [:]
+        for oid in oids {
+            names[oid] = unresolved
+        }
+        for row in rows.compactMap(row(fromColumns:)) {
+            names[row.oid] = typeName(for: row)
+        }
+        return names
+    }
+}

Plugins/PostgreSQLDriverPlugin/PostgreSQLPluginDriver+EnumTypes.swift

@@ -16,17 +16,8 @@ extension PostgreSQLPluginDriver {
     func probeEnumOids() async {
         do {
             let result = try await core.execute(query: PostgreSQLSchemaQueries.enumTypeOidQuery)
-            var map: [UInt32: String] = [:]
-            for row in result.rows {
-                guard row.count >= 3,
-                      let typeName = row[2].asText,
-                      let scalarOid = row[0].asText.flatMap({ UInt32($0) }) else { continue }
-                map[scalarOid] = "ENUM(\(typeName))"
-                if let arrayOid = row[1].asText.flatMap({ UInt32($0) }), arrayOid != 0 {
-                    map[arrayOid] = "ENUM[](\(typeName))"
-                }
-            }
-            core.setEnumOidMap(map)
+            let rows = result.rows.map { row in row.map(\.asText) }
+            core.mergeCatalogTypeNames(PostgreSQLCatalogTypeNames.enumProbeNames(rows: rows))
         } catch {
             enumProbeLogger.debug(
                 "Enum OID probe failed; enum columns fall back to text for this session: \(error.localizedDescription)"