fix(plugin-postgresql): keep the selected schema applied after reconnect

CHANGELOG.md

@@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- PostgreSQL: the selected schema now stays applied to editor queries after an automatic reconnect, so unqualified table names keep resolving against it instead of falling back to the default schema. (#1540)
 - Import now detects the Setapp edition of TablePlus and reads connections from its data folder. It was reported as not installed before. (#1528)
 - Favorite keyword suggestions now show in the editor autocomplete when you type the keyword. They were being dropped before reaching the popup.
 

Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift

@@ -14,6 +14,7 @@ final class LibPQDriverCore: @unchecked Sendable {
     private var libpqConnection: LibPQPluginConnection?
 
     var currentSchema: String = "public"
+    private var selectedSchema: String?
 
     var onPostConnect: (@Sendable () async -> Void)?
 
@@ -45,9 +46,20 @@ final class LibPQDriverCore: @unchecked Sendable {
             currentSchema = schema
         }
 
+        if let selectedSchema,
+           (try? await pqConn.executeQuery(PostgreSQLSchemaQueries.setSearchPath(toSchema: selectedSchema))) != nil {
+            currentSchema = selectedSchema
+        }
+
         await onPostConnect?()
     }
 
+    func applySchema(_ schema: String) async throws {
+        _ = try await execute(query: PostgreSQLSchemaQueries.setSearchPath(toSchema: schema))
+        selectedSchema = schema
+        currentSchema = schema
+    }
+
     func disconnect() {
         libpqConnection?.disconnect()
         libpqConnection = nil
@@ -180,9 +192,7 @@ extension LibPQBackedDriver {
     }
 
     func switchSchema(to schema: String) async throws {
-        let escapedName = schema.replacingOccurrences(of: "\"", with: "\"\"")
-        _ = try await core.execute(query: "SET search_path TO \"\(escapedName)\", public")
-        core.currentSchema = schema
+        try await core.applySchema(schema)
     }
 
     var currentSchema: String? { core.currentSchema }

Plugins/PostgreSQLDriverPlugin/PostgreSQLSchemaQueries.swift

@@ -75,4 +75,9 @@ enum PostgreSQLSchemaQueries {
 
         return unions.joined(separator: "\nUNION ALL\n") + "\nORDER BY table_name"
     }
+
+    static func setSearchPath(toSchema schema: String) -> String {
+        let quotedIdentifier = schema.replacingOccurrences(of: "\"", with: "\"\"")
+        return "SET search_path TO \"\(quotedIdentifier)\", public"
+    }
 }

TableProTests/Plugins/PostgreSQLSearchPathTests.swift

@@ -0,0 +1,43 @@
+//
+//  PostgreSQLSearchPathTests.swift
+//  TableProTests
+//
+
+import Foundation
+import Testing
+
+@Suite("PostgreSQLSchemaQueries.setSearchPath")
+struct PostgreSQLSearchPathTests {
+    @Test("quotes the schema as an identifier and keeps public on the path")
+    func plainSchema() {
+        #expect(
+            PostgreSQLSchemaQueries.setSearchPath(toSchema: "analytics")
+                == "SET search_path TO \"analytics\", public"
+        )
+    }
+
+    @Test("preserves mixed-case schema names with identifier quoting")
+    func mixedCaseSchema() {
+        #expect(
+            PostgreSQLSchemaQueries.setSearchPath(toSchema: "MySchema")
+                == "SET search_path TO \"MySchema\", public"
+        )
+    }
+
+    @Test("doubles embedded double quotes so the name stays a single identifier")
+    func schemaWithEmbeddedQuote() {
+        #expect(
+            PostgreSQLSchemaQueries.setSearchPath(toSchema: "wei\"rd")
+                == "SET search_path TO \"wei\"\"rd\", public"
+        )
+    }
+
+    @Test("neutralizes a quote-break injection attempt")
+    func injectionAttempt() {
+        let malicious = "public\"; DROP TABLE users; --"
+        #expect(
+            PostgreSQLSchemaQueries.setSearchPath(toSchema: malicious)
+                == "SET search_path TO \"public\"\"; DROP TABLE users; --\", public"
+        )
+    }
+}