diff --git a/CHANGELOG.md b/CHANGELOG.md index 697487d63..4c748b460 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift b/Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift index 2389abf18..5102d61b3 100644 --- a/Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift +++ b/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 } diff --git a/Plugins/PostgreSQLDriverPlugin/PostgreSQLSchemaQueries.swift b/Plugins/PostgreSQLDriverPlugin/PostgreSQLSchemaQueries.swift index 1b99977f2..8a5cd6332 100644 --- a/Plugins/PostgreSQLDriverPlugin/PostgreSQLSchemaQueries.swift +++ b/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" + } } diff --git a/TableProTests/Plugins/PostgreSQLSearchPathTests.swift b/TableProTests/Plugins/PostgreSQLSearchPathTests.swift new file mode 100644 index 000000000..d6dbc0756 --- /dev/null +++ b/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" + ) + } +}