feat(kingbasees): add a registry driver

.github/plugin-registry.json

@@ -43,6 +43,9 @@
     "KafkaDriverPlugin": {
       "name": "Kafka Driver"
     },
+    "KingbaseESDriverPlugin": {
+      "name": "KingbaseES Driver"
+    },
     "LibSQLDriverPlugin": {
       "name": "libSQL / Turso Driver"
     },

apple/Packages/TableProCore/Sources/TableProSQLGrammar/PLSQLUnitTracker.swift

@@ -184,6 +184,17 @@ public struct PLSQLUnitTracker: SQLStatementBoundaryTracking {
         settlePendingBeforeNonWord()
     }
 
+    /// KingbaseES in Oracle mode also takes PostgreSQL's quoted routine body, `AS 'SELECT 1'` or `AS $$ ... $$`, which
+    /// opens no block. Oracle never puts a literal after a subprogram's `IS` or `AS`; a quoted name there still starts
+    /// its declarations.
+    public mutating func observeQuotedToken(opener: UInt16) {
+        if pendingBodyKeyword, opener != SqlLexer.doubleQuote {
+            pendingBodyKeyword = false
+            header = nil
+        }
+        observeOpaqueToken()
+    }
+
     /// PL/SQL reads `a . b` as `a.b`, so whitespace leaves a member name a member name.
     public mutating func observeGap() {}
 

apple/Packages/TableProCore/Sources/TableProSQLGrammar/SQLLexicalProfile.swift

@@ -219,6 +219,13 @@ public struct SQLLexicalProfile: Sendable, Hashable {
             "AlloyDB": postgreSQLFamily,
             "Citus": postgreSQLFamily,
             "PGlite": postgreSQLFamily,
+            "KingbaseES": SQLLexicalProfile(
+                grammar: postgreSQL,
+                undetermined: [
+                    .backslashEscapesInSingleQuotes, .alternativeQuoting, .plsqlBlocks,
+                    .slashLineTerminators, .dollarAndHashInIdentifiers,
+                ]
+            ),
             "Redshift": SQLLexicalProfile(
                 grammar: postgreSQL,
                 undetermined: [.backslashEscapesInSingleQuotes, .nestedBlockComments, .escapeStringPrefix]

apple/Packages/TableProCore/Sources/TableProSQLGrammar/SQLStatementBoundaryTracking.swift

@@ -20,7 +20,8 @@ public enum SQLStatementTerminator: Equatable, Sendable {
 /// boundaries without looking ahead.
 ///
 /// Words arrive uppercased. Whitespace and comments arrive as ``observeGap()``; a string literal, a quoted identifier
-/// or a dollar-quoted body arrives as ``observeOpaqueToken()``.
+/// or a dollar-quoted body arrives as ``observeQuotedToken(opener:)``, and any other token no word or symbol describes
+/// as ``observeOpaqueToken()``.
 public protocol SQLStatementBoundaryTracking {
     /// Whether words still bear on where this statement ends. A reader that has to assemble words itself can stop
     /// doing so once this is false, which keeps a plain `INSERT` dump free of the cost.
@@ -36,6 +37,10 @@ public protocol SQLStatementBoundaryTracking {
     mutating func observeSymbol(_ symbol: UInt16)
     mutating func observeOpaqueToken()
 
+    /// A quoted token, named by its first character: `'`, `$`, or a prefix such as `E` or `q` for a literal, and `"`,
+    /// a backtick or `[` for a quoted name.
+    mutating func observeQuotedToken(opener: UInt16)
+
     /// Whitespace or a comment between two tokens, which keeps a word from continuing the name before it: `x$MERGE`
     /// is one name in T-SQL, while `x$` and `MERGE` on the next line are a name and a statement.
     mutating func observeGap()
@@ -46,6 +51,12 @@ public protocol SQLStatementBoundaryTracking {
     mutating func reset()
 }
 
+public extension SQLStatementBoundaryTracking {
+    mutating func observeQuotedToken(opener: UInt16) {
+        observeOpaqueToken()
+    }
+}
+
 public enum SQLStatementBoundaries {
     /// The one place a grammar is matched to its statement boundaries, so no reader can pick a different one.
     public static func makeTracker(for grammar: SQLLexicalGrammar) -> any SQLStatementBoundaryTracking {

apple/Packages/TableProCore/Sources/TableProSQLGrammar/SQLStatementScanner.swift

@@ -333,7 +333,10 @@ public enum SQLStatementScanner {
                     tracker.observeGap()
                 case .executableComment:
                     hasStatementContent = true
-                case .quoted, .parameter:
+                case .quoted:
+                    hasStatementContent = true
+                    tracker.observeQuotedToken(opener: nsQuery.character(at: span.start))
+                case .parameter:
                     hasStatementContent = true
                     tracker.observeOpaqueToken()
                 }

apple/Packages/TableProCore/Sources/TableProSQLVocabulary/SQLEngineVocabulary.swift

@@ -67,6 +67,7 @@ public struct SQLEngineVocabulary: Sendable, Hashable {
         "PostgreSQL": postgreSQL,
         "CockroachDB": postgreSQL,
         "PGlite": postgreSQL,
+        "KingbaseES": postgreSQL,
         "Redshift": redshift,
         "SQLite": sqlite,
         "SQL Server": sqlServer,

apple/Packages/TableProCore/Tests/TableProPluginKitTests/KingbaseESDialectTests.swift

@@ -0,0 +1,13 @@
+import TableProPluginKit
+import Testing
+
+struct KingbaseESDialectTests {
+    @Test("KingbaseES keeps PostgreSQL dollar quotes and escape strings")
+    func postgresLexing() {
+        let dialect = SqlDialect.from(databaseTypeId: "KingbaseES")
+        #expect(dialect == .postgres)
+        #expect(dialect.supportsDollarQuotes)
+        #expect(dialect.supportsEscapeStringPrefix)
+        #expect(!dialect.requiresBackslashEscapesInSingleQuotes)
+    }
+}

apple/Packages/TableProCore/Tests/TableProSQLGrammarTests/SQLLexicalReadingsTests.swift

@@ -7,7 +7,7 @@ struct SQLLexicalReadingsTests {
     @Test("Every SQL engine TablePro ships has a curated grammar")
     func curatedTableCoversShippedEngines() {
         let shipped = [
-            "MySQL", "MariaDB", "TiDB", "Databend", "OceanBase", "PostgreSQL", "Redshift", "CockroachDB", "PGlite",
+            "MySQL", "MariaDB", "TiDB", "Databend", "OceanBase", "PostgreSQL", "Redshift", "CockroachDB", "PGlite", "KingbaseES",
             "SQLite", "libSQL", "Turso", "Cloudflare D1", "DuckDB", "Oracle", "Dameng", "SQL Server", "ClickHouse",
             "Snowflake", "BigQuery", "Spanner", "Trino", "Presto", "Teradata", "SAP HANA", "Cassandra", "ScyllaDB", "DynamoDB",
             "SurrealDB", "Redis", "MongoDB", "etcd", "Elasticsearch", "Typesense", "Weaviate", "Kafka",
@@ -130,4 +130,35 @@ struct SQLLexicalReadingsTests {
         ])
         #expect(profile.readings.first == .ansi)
     }
+    @Test("KingbaseES retains both compatibility-mode readings before a session is known")
+    func kingbaseCompatibilityReadings() throws {
+        let profile = try #require(SQLLexicalProfile.curated(forDatabaseTypeId: "KingbaseES"))
+        #expect(profile.grammar.contains(.taggedDollarQuotes))
+        #expect(profile.readings.contains { $0.contains(.alternativeQuoting) && $0.contains(.plsqlBlocks) })
+        let oracle = SQLSessionLexicalFacts(
+            determined: [.alternativeQuoting, .plsqlBlocks], enabled: [.alternativeQuoting, .plsqlBlocks]
+        )
+        let readings = SQLLexicalReadings.resolve(databaseTypeId: "KingbaseES", declared: nil, session: oracle)
+        #expect(readings.execution.contains(.alternativeQuoting))
+        #expect(readings.execution.contains(.plsqlBlocks))
+        #expect(readings.all.contains { !$0.contains(.alternativeQuoting) })
+    }
+
+    // Measured on KingbaseES V9R1C10 in Oracle mode: each of these creates its function, beside PL/SQL bodies.
+    @Test("Oracle-mode KingbaseES ends a quoted routine body at its semicolon and keeps a PL/SQL body whole")
+    func kingbaseOracleModeRoutineBodies() {
+        let oracle = SQLSessionLexicalFacts(
+            determined: [.alternativeQuoting, .plsqlBlocks, .slashLineTerminators, .dollarAndHashInIdentifiers],
+            enabled: [.alternativeQuoting, .plsqlBlocks, .slashLineTerminators, .dollarAndHashInIdentifiers]
+        )
+        let grammar = SQLLexicalReadings.resolve(databaseTypeId: "KingbaseES", declared: nil, session: oracle).execution
+        func sent(_ sql: String) -> [String] {
+            SQLStatementScanner.executableStatements(in: sql, grammar: grammar).map(\.sql)
+        }
+        #expect(sent("CREATE FUNCTION f() RETURNS integer LANGUAGE SQL AS 'SELECT 42'; SELECT 2;").count == 2)
+        #expect(sent("CREATE FUNCTION f() RETURNS integer AS $$ SELECT 1; $$ LANGUAGE sql; SELECT 2;").count == 2)
+        #expect(sent("CREATE FUNCTION f() RETURNS integer AS $b$ BEGIN RETURN 1; END; $b$ LANGUAGE plpgsql; SELECT 2;").count == 2)
+        #expect(sent("CREATE FUNCTION f RETURN NUMBER AS BEGIN RETURN 1; END;\n/\nSELECT 2;").count == 2)
+        #expect(sent("CREATE FUNCTION f RETURN NUMBER AS \"v\" NUMBER := 1; BEGIN RETURN \"v\"; END;\n/\nSELECT 2;").count == 2)
+    }
 }

apple/Packages/TableProCore/Tests/TableProSQLVocabularyTests/SQLEngineVocabularyTests.swift

@@ -9,7 +9,7 @@ struct SQLEngineVocabularyTests {
         let typeIds = SQLEngineVocabulary.curatedDatabaseTypeIds
         #expect(typeIds == [
             "MySQL", "MariaDB", "TiDB", "Databend", "OceanBase",
-            "PostgreSQL", "CockroachDB", "PGlite", "Redshift",
+            "PostgreSQL", "CockroachDB", "PGlite", "KingbaseES", "Redshift",
             "SQLite", "SQL Server", "Oracle", "DuckDB",
         ])
         for typeId in typeIds {
@@ -32,6 +32,7 @@ struct SQLEngineVocabularyTests {
         ("OceanBase", "MySQL"),
         ("CockroachDB", "PostgreSQL"),
         ("PGlite", "PostgreSQL"),
+        ("KingbaseES", "PostgreSQL"),
     ])
     func variantResolvesLikeItsBase(variant: String, base: String) {
         let resolved = SQLEngineVocabulary.curated(forDatabaseTypeId: variant)

apple/Packages/TableProMacCore/Sources/TableProUtilities/SQL/Folding/SQLFoldScanner.swift

@@ -243,10 +243,13 @@ private struct Scan {
             appendSpanningRegion(.blockComment, start: start, startLine: startLine, end: span.contentEnd)
         case .quoted where text.character(at: index) == SqlDollarQuote.dollar:
             openStatementIfNeeded()
-            boundaries.observeOpaqueToken()
+            boundaries.observeQuotedToken(opener: SqlDollarQuote.dollar)
             line += span.newlines
             appendSpanningRegion(.quotedBody, start: start, startLine: startLine, end: span.contentEnd)
-        case .quoted, .parameter:
+        case .quoted:
+            boundaries.observeQuotedToken(opener: text.character(at: index))
+            line += span.newlines
+        case .parameter:
             boundaries.observeOpaqueToken()
             line += span.newlines
         }

apple/Packages/TableProMacCore/Sources/TableProUtilities/SQL/SQLStateChangingFunctions.swift

@@ -2,7 +2,10 @@ import Foundation
 
 public enum SQLStateChangingFunctions {
     public static func contains(_ name: String) -> Bool {
-        names.contains(name.lowercased())
+        let lowered = name.lowercased()
+        // KingbaseES ships each pg_ function under a sys_ name too, such as sys_terminate_backend (measured on V9R1C10).
+        if lowered.hasPrefix("sys_") { return names.contains("pg_" + lowered.dropFirst(4)) }
+        return names.contains(lowered)
     }
 
     private static let names: Set<String> = serverSignals

apple/Plugins/FormatPluginTests/SQLExportSessionDeclarationTests.swift

@@ -268,4 +268,11 @@ struct SQLExportSessionDeclarationTests {
         #expect(parts.count == 1)
         #expect(!writer.didSplit)
     }
+    @Test("KingbaseES dumps declare PostgreSQL string and encoding semantics")
+    func kingbaseSessionDeclaration() {
+        #expect(SQLExportSessionDeclaration.forDatabaseType("KingbaseES")
+            == SQLExportSessionDeclaration.forDatabaseType("PostgreSQL"))
+        #expect(!SQLExportSessionDeclaration.forDatabaseType("KingbaseES").prologue.isEmpty)
+    }
+
 }

apple/Plugins/KingbaseESDriverPlugin/Info.plist

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>TableProPluginKitVersion</key>
+	<integer>TABLEPRO_PLUGINKIT_VERSION</integer>
+	<key>TableProProvidesDatabaseTypeIds</key>
+	<array>
+		<string>KingbaseES</string>
+	</array>
+</dict>
+</plist>

apple/Plugins/KingbaseESDriverPlugin/KingbaseESCompatibilityMode.swift

@@ -0,0 +1,37 @@
+import Foundation
+import TableProPluginKit
+
+/// The SQL a KingbaseES server speaks, fixed when its cluster is initialized.
+enum KingbaseESCompatibilityMode: Equatable, Sendable {
+    case postgreSQL
+    case oracle
+
+    // The two-argument form answers NULL where the parameter does not exist, which reads as PostgreSQL mode.
+    static let probeQuery = "SELECT current_setting('database_mode', true)"
+
+    static let lexicalFeaturesItDecides: SQLLexicalFeatures = [
+        .alternativeQuoting, .plsqlBlocks, .slashLineTerminators, .dollarAndHashInIdentifiers,
+    ]
+
+    init(setting: String?) throws {
+        switch setting?.lowercased() {
+        case nil, "pg":
+            self = .postgreSQL
+        case "oracle":
+            self = .oracle
+        case let other?:
+            throw LibPQPluginError(
+                message: String(
+                    format: String(localized: "KingbaseES in %@ mode is not supported. Connect to a server in PostgreSQL or Oracle mode."),
+                    other
+                ),
+                sqlState: nil,
+                detail: nil
+            )
+        }
+    }
+
+    var lexicalFeatures: SQLLexicalFeatures {
+        self == .oracle ? Self.lexicalFeaturesItDecides : []
+    }
+}

apple/Plugins/KingbaseESDriverPlugin/KingbaseESPlugin.swift

@@ -0,0 +1,72 @@
+import Foundation
+import TableProPluginKit
+
+final class KingbaseESPlugin: NSObject, TableProPlugin, DriverPlugin {
+    static let pluginName = "KingbaseES Driver"
+    static let pluginVersion = "1.0.0"
+    static let pluginDescription = "KingbaseES database support via libpq"
+    static let capabilities: [PluginCapability] = [.databaseDriver]
+
+    static let databaseTypeId = "KingbaseES"
+    static let databaseDisplayName = "KingbaseES"
+    static let iconName = "cylinder"
+    static let defaultPort = 54_321
+    static let isDownloadable = true
+    static let urlSchemes = ["kingbase", "kingbase8"]
+    static let brandColorHex = "#C9141E"
+    static let systemDatabaseNames = PostgreSQLPlugin.systemDatabaseNames
+    static let systemSchemaNames = [
+        "pg_catalog", "information_schema", "pg_toast", "sys_catalog", "sys", "anon", "dbms_job", "dbms_scheduler",
+        "dbms_sql", "kdb_schedule", "perf", "pg_bitmapindex", "src_restrict", "sys_hm", "sysaudit", "sysmac", "wmsys",
+        "xlog_record_read"
+    ]
+
+    static let additionalConnectionFields: [ConnectionField] = [
+        ConnectionField(
+            id: "usePgpass",
+            label: String(localized: "Use ~/.pgpass"),
+            defaultValue: "false",
+            fieldType: .toggle,
+            section: .authentication,
+            hidesPassword: true
+        ),
+        ConnectionField(
+            id: "connectionOptions",
+            label: String(localized: "Connection Options"),
+            fieldType: .text,
+            section: .advanced
+        )
+    ]
+
+    static let supportsSchemaSwitching = PostgreSQLPlugin.supportsSchemaSwitching
+    static let postConnectActions = PostgreSQLPlugin.postConnectActions
+    static let explainVariants = PostgreSQLPlugin.explainVariants
+    static let databaseGroupingStrategy = PostgreSQLPlugin.databaseGroupingStrategy
+    static let columnTypesByCategory = PostgreSQLPlugin.columnTypesByCategory
+    static let supportsCascadeDrop = PostgreSQLPlugin.supportsCascadeDrop
+    static let supportsForeignKeyDisable = PostgreSQLPlugin.supportsForeignKeyDisable
+    static let requiresReconnectForDatabaseSwitch = PostgreSQLPlugin.requiresReconnectForDatabaseSwitch
+    static let parameterStyle = PostgreSQLPlugin.parameterStyle
+    static let supportsDropDatabase = PostgreSQLPlugin.supportsDropDatabase
+    static let supportsRenameTable = PostgreSQLPlugin.supportsRenameTable
+    static let supportsRenameDatabase = PostgreSQLPlugin.supportsRenameDatabase
+    static let supportsRenameSchema = PostgreSQLPlugin.supportsRenameSchema
+    static let supportsDropSchema = PostgreSQLPlugin.supportsDropSchema
+    static let supportsCreateSchema = PostgreSQLPlugin.supportsCreateSchema
+    static let supportsSchemaOwner = PostgreSQLPlugin.supportsSchemaOwner
+    static let supportsSchemaPrivileges = PostgreSQLPlugin.supportsSchemaPrivileges
+    static let supportsTriggers = PostgreSQLPlugin.supportsTriggers
+    static let supportsRoutines = PostgreSQLPlugin.supportsRoutines
+    static let supportsDatabaseTriggerBrowse = PostgreSQLPlugin.supportsDatabaseTriggerBrowse
+    static let supportsUserDefinedTypeBrowse = PostgreSQLPlugin.supportsUserDefinedTypeBrowse
+    static let supportsTriggerEditing = PostgreSQLPlugin.supportsTriggerEditing
+    static let structureColumnFields = PostgreSQLPlugin.structureColumnFields
+    static let supportsCheckConstraints = PostgreSQLPlugin.supportsCheckConstraints
+    static let supportsCheckConstraintEditing = PostgreSQLPlugin.supportsCheckConstraintEditing
+    static let supportsGeneratedColumns = PostgreSQLPlugin.supportsGeneratedColumns
+    static let sqlDialect = PostgreSQLPlugin.sqlDialect
+
+    func createDriver(config: DriverConnectionConfig) -> any PluginDatabaseDriver {
+        KingbaseESPluginDriver(config: config)
+    }
+}

apple/Plugins/KingbaseESDriverPlugin/KingbaseESPluginDriver.swift

@@ -0,0 +1,23 @@
+import Foundation
+import os
+import TableProPluginKit
+
+final class KingbaseESPluginDriver: PostgreSQLPluginDriver, @unchecked Sendable {
+    private let mode = OSAllocatedUnfairLock(initialState: KingbaseESCompatibilityMode.postgreSQL)
+
+    override var sessionLexicalState: PluginSessionLexicalState? {
+        guard let inherited = super.sessionLexicalState else { return nil }
+        let modeFeatures = mode.withLock { $0 }.lexicalFeatures
+        return PluginSessionLexicalState(
+            determined: inherited.determined.union(KingbaseESCompatibilityMode.lexicalFeaturesItDecides),
+            enabled: inherited.enabled.union(modeFeatures)
+        )
+    }
+
+    override func probeAfterConnect() async throws {
+        let result = try await execute(query: KingbaseESCompatibilityMode.probeQuery)
+        let detected = try KingbaseESCompatibilityMode(setting: result.rows.first?.first?.asText)
+        mode.withLock { $0 = detected }
+        try await super.probeAfterConnect()
+    }
+}

apple/Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift

@@ -25,7 +25,7 @@ final class LibPQDriverCore: @unchecked Sendable {
     var currentSchema: String = "public"
     private var selectedSchema: String?
 
-    var onPostConnect: (@Sendable () async -> Void)?
+    var onPostConnect: (@Sendable () async throws -> Void)?
 
     /// Set by `LibPQBackedDriver` for the span of one connect, so every driver built on this
     /// core reports its handshake steps without having to thread a parameter through its own
@@ -124,7 +124,7 @@ final class LibPQDriverCore: @unchecked Sendable {
                     currentSchema = selectedSchema
                 }
 
-                await onPostConnect?()
+                try await onPostConnect?()
                 try await pqConn.finishConnecting()
                 try Task.checkCancellation()
             } onCancel: {