Skip to content

Commit afc0529

Browse files
authored
Merge pull request #141 from datlechin/fix/mongodb-database-switch-and-parser
fix: MongoDB database switch, dotted collection names, and shell parser
2 parents 4a85b97 + 1092027 commit afc0529

11 files changed

Lines changed: 831 additions & 29 deletions

TablePro/Core/Database/MongoDBDriver.swift

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import OSLog
1313
/// Parses mongo shell syntax (db.collection.find/insert/update/delete)
1414
/// and dispatches to MongoDBConnection for execution.
1515
final class MongoDBDriver: DatabaseDriver {
16-
let connection: DatabaseConnection
16+
private(set) var connection: DatabaseConnection
1717
private(set) var status: ConnectionStatus = .disconnected
1818

1919
private var mongoConnection: MongoDBConnection?
@@ -24,6 +24,10 @@ final class MongoDBDriver: DatabaseDriver {
2424
self.connection = connection
2525
}
2626

27+
func switchDatabase(to database: String) {
28+
connection.database = database
29+
}
30+
2731
// MARK: - Server Version
2832

2933
var serverVersion: String? {
@@ -408,7 +412,8 @@ final class MongoDBDriver: DatabaseDriver {
408412
opts.append("\"name\": \"\(name)\"")
409413

410414
let optsJson = "{\(opts.joined(separator: ", "))}"
411-
sections.append("db.\(table).createIndex(\(keyJson), \(optsJson))")
415+
let escapedTable = table.replacingOccurrences(of: "\\", with: "\\\\").replacingOccurrences(of: "\"", with: "\\\"")
416+
sections.append("db[\"\(escapedTable)\"].createIndex(\(keyJson), \(optsJson))")
412417
}
413418
}
414419
} catch {

TablePro/Core/MongoDB/MongoDBStatementGenerator.swift

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@ struct MongoDBStatementGenerator {
1515
let collectionName: String
1616
let columns: [String]
1717

18+
/// Collection accessor using bracket notation for safety with dotted names
19+
private var collectionAccessor: String {
20+
"db[\"\(escapeJsonString(collectionName))\"]"
21+
}
22+
1823
/// Index of "_id" field in the columns array (used as primary key equivalent)
1924
var idColumnIndex: Int? {
2025
columns.firstIndex(of: "_id")
@@ -97,7 +102,7 @@ struct MongoDBStatementGenerator {
97102
guard !doc.isEmpty else { return nil }
98103

99104
let docJson = serializeDocument(doc)
100-
let shell = "db.\(collectionName).insertOne(\(docJson))"
105+
let shell = "\(collectionAccessor).insertOne(\(docJson))"
101106
return ParameterizedStatement(sql: shell, parameters: [])
102107
}
103108

@@ -134,7 +139,7 @@ struct MongoDBStatementGenerator {
134139
guard docs.count > 1 else { return nil }
135140

136141
let docsArray = "[\(docs.joined(separator: ", "))]"
137-
let shell = "db.\(collectionName).insertMany(\(docsArray))"
142+
let shell = "\(collectionAccessor).insertMany(\(docsArray))"
138143
return ParameterizedStatement(sql: shell, parameters: [])
139144
}
140145

@@ -179,7 +184,7 @@ struct MongoDBStatementGenerator {
179184
}
180185

181186
let updateJson = "{\(updateParts.joined(separator: ", "))}"
182-
let shell = "db.\(collectionName).updateOne(\(filterJson), \(updateJson))"
187+
let shell = "\(collectionAccessor).updateOne(\(filterJson), \(updateJson))"
183188
return ParameterizedStatement(sql: shell, parameters: [])
184189
}
185190

@@ -206,7 +211,7 @@ struct MongoDBStatementGenerator {
206211
}
207212

208213
let inList = idValues.joined(separator: ", ")
209-
let shell = "db.\(collectionName).deleteMany({\"_id\": {\"$in\": [\(inList)]}})"
214+
let shell = "\(collectionAccessor).deleteMany({\"_id\": {\"$in\": [\(inList)]}})"
210215
return ParameterizedStatement(sql: shell, parameters: [])
211216
}
212217

@@ -220,7 +225,7 @@ struct MongoDBStatementGenerator {
220225
idIndex < originalRow.count,
221226
let idValue = originalRow[idIndex] {
222227
let filterJson = buildIdFilter(idValue)
223-
let shell = "db.\(collectionName).deleteOne(\(filterJson))"
228+
let shell = "\(collectionAccessor).deleteOne(\(filterJson))"
224229
return ParameterizedStatement(sql: shell, parameters: [])
225230
}
226231

@@ -236,7 +241,7 @@ struct MongoDBStatementGenerator {
236241
guard !filter.isEmpty else { return nil }
237242

238243
let filterJson = serializeDocument(filter)
239-
let shell = "db.\(collectionName).deleteOne(\(filterJson))"
244+
let shell = "\(collectionAccessor).deleteOne(\(filterJson))"
240245
return ParameterizedStatement(sql: shell, parameters: [])
241246
}
242247

TablePro/Core/MongoDB/MongoShellParser.swift

Lines changed: 120 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,11 @@ struct MongoShellParser {
100100
return .runCommand(command: arg)
101101
}
102102

103+
// db["collection"].method(args) bracket notation
104+
if trimmed.hasPrefix("db[") {
105+
return try parseBracketExpression(trimmed)
106+
}
107+
103108
// db.collection.method(args) pattern
104109
guard trimmed.hasPrefix("db.") else {
105110
throw MongoShellParseError.invalidSyntax("Query must start with 'db.' or be a JSON command")
@@ -110,23 +115,133 @@ struct MongoShellParser {
110115

111116
// MARK: - Private Parsing
112117

118+
/// Parse db["collection"].method(args) bracket notation.
119+
/// Supports both double and single quotes around the collection name.
120+
private static func parseBracketExpression(_ input: String) throws -> MongoOperation {
121+
// input starts with db[
122+
let afterBracket = String(input.dropFirst(3)) // drop "db["
123+
124+
// Determine quote character (" or ')
125+
guard let quoteChar = afterBracket.first, quoteChar == "\"" || quoteChar == "'" else {
126+
throw MongoShellParseError.invalidSyntax("Expected quoted collection name in db[...]")
127+
}
128+
129+
// Find closing quote (handle escaped quotes)
130+
var collectionName = ""
131+
var i = afterBracket.index(after: afterBracket.startIndex)
132+
var escapeNext = false
133+
while i < afterBracket.endIndex {
134+
let ch = afterBracket[i]
135+
if escapeNext {
136+
collectionName.append(ch)
137+
escapeNext = false
138+
i = afterBracket.index(after: i)
139+
continue
140+
}
141+
if ch == "\\" {
142+
escapeNext = true
143+
i = afterBracket.index(after: i)
144+
continue
145+
}
146+
if ch == quoteChar {
147+
break
148+
}
149+
collectionName.append(ch)
150+
i = afterBracket.index(after: i)
151+
}
152+
153+
guard i < afterBracket.endIndex else {
154+
throw MongoShellParseError.invalidSyntax("Unterminated string in db[...]")
155+
}
156+
157+
// Move past closing quote and expect "]"
158+
i = afterBracket.index(after: i)
159+
guard i < afterBracket.endIndex, afterBracket[i] == "]" else {
160+
throw MongoShellParseError.invalidSyntax("Expected ']' after collection name in db[...]")
161+
}
162+
i = afterBracket.index(after: i)
163+
164+
let remaining = String(afterBracket[i...]).trimmingCharacters(in: .whitespacesAndNewlines)
165+
166+
// No method chain — treat as find all
167+
if remaining.isEmpty {
168+
return .find(collection: collectionName, filter: "{}", options: MongoFindOptions())
169+
}
170+
171+
// Expect ".method(args)" after db["collection"]
172+
guard remaining.hasPrefix(".") else {
173+
throw MongoShellParseError.invalidSyntax("Expected '.method()' after db[\"...\"]")
174+
}
175+
176+
let methodChain = String(remaining.dropFirst())
177+
return try parseMethodChain(collection: collectionName, chain: methodChain)
178+
}
179+
113180
private static func parseDbExpression(_ input: String) throws -> MongoOperation {
114181
// Remove "db." prefix
115182
let afterDb = String(input.dropFirst(3))
116183

117-
// Find the collection name (everything before the first ".")
118-
guard let dotIndex = afterDb.firstIndex(of: ".") else {
119-
// Just "db.collectionName" -- treat as find all
184+
guard let firstParen = afterDb.firstIndex(of: "(") else {
185+
// No parentheses at all — "db.collectionName" or "db.system.version" — treat as find all
120186
let collection = afterDb.trimmingCharacters(in: .whitespacesAndNewlines)
187+
guard !collection.isEmpty else {
188+
throw MongoShellParseError.invalidSyntax("Missing collection name after 'db.'")
189+
}
121190
return .find(collection: collection, filter: "{}", options: MongoFindOptions())
122191
}
123192

124-
let collection = String(afterDb[afterDb.startIndex..<dotIndex])
125-
let remainder = String(afterDb[afterDb.index(after: dotIndex)...])
193+
// Find the last "." before the first "(". Everything before it is the collection name,
194+
// and everything from it onward is the method chain.
195+
// This correctly handles dotted collection names like "system.version".
196+
let beforeParen = afterDb[afterDb.startIndex..<firstParen]
197+
guard let lastDot = beforeParen.lastIndex(of: ".") else {
198+
// No dot before paren — db-level method call like db.getCollectionNames()
199+
return try parseDbLevelMethod(afterDb)
200+
}
201+
202+
let collection = String(afterDb[afterDb.startIndex..<lastDot])
203+
let remainder = String(afterDb[afterDb.index(after: lastDot)...])
126204

127205
return try parseMethodChain(collection: collection, chain: remainder)
128206
}
129207

208+
/// Parse a db-level method call like db.getCollectionNames(), db.stats(), etc.
209+
/// Input is the string after "db." — e.g. "getCollectionNames()" or "createCollection(\"test\")"
210+
private static func parseDbLevelMethod(_ input: String) throws -> MongoOperation {
211+
guard let parenIndex = input.firstIndex(of: "(") else {
212+
throw MongoShellParseError.invalidSyntax("Expected method call with parentheses")
213+
}
214+
215+
let methodName = String(input[input.startIndex..<parenIndex])
216+
let argAndRest = try extractParenthesizedArgAndRemainder(from: input, startingAt: parenIndex)
217+
let arg = argAndRest.arg
218+
219+
switch methodName {
220+
case "getCollectionNames", "listCollections":
221+
return .listCollections
222+
223+
case "createCollection":
224+
let name = arg.trimmingCharacters(in: .whitespacesAndNewlines)
225+
.trimmingCharacters(in: CharacterSet(charactersIn: "\"'"))
226+
guard !name.isEmpty else {
227+
throw MongoShellParseError.missingArgument("createCollection requires a collection name")
228+
}
229+
return .runCommand(command: "{ \"create\": \"\(name)\" }")
230+
231+
case "dropDatabase":
232+
return .runCommand(command: "{ \"dropDatabase\": 1 }")
233+
234+
case "version":
235+
return .runCommand(command: "{ \"buildInfo\": 1 }")
236+
237+
case "stats":
238+
return .runCommand(command: "{ \"dbStats\": 1 }")
239+
240+
default:
241+
throw MongoShellParseError.unsupportedMethod(methodName)
242+
}
243+
}
244+
130245
private static func parseMethodChain(collection: String, chain: String) throws -> MongoOperation {
131246
guard let parenIndex = chain.firstIndex(of: "(") else {
132247
throw MongoShellParseError.invalidSyntax("Expected method call with parentheses")

TablePro/Core/Services/ExportService.swift

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1159,7 +1159,8 @@ final class ExportService: ObservableObject {
11591159
}
11601160
if foundHeader {
11611161
var processedLine = line
1162-
let ddlAccessor = "db.\(collection)"
1162+
let escapedForDDL = collection.replacingOccurrences(of: "\\", with: "\\\\").replacingOccurrences(of: "\"", with: "\\\"")
1163+
let ddlAccessor = "db[\"\(escapedForDDL)\"]"
11631164
if processedLine.hasPrefix(ddlAccessor) {
11641165
processedLine = collectionAccessor + processedLine.dropFirst(ddlAccessor.count)
11651166
}

TablePro/Models/QueryTab.swift

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -565,7 +565,8 @@ final class QueryTabManager: ObservableObject {
565565
let pageSize = AppSettingsManager.shared.dataGrid.defaultPageSize
566566
let query: String
567567
if databaseType == .mongodb {
568-
query = "db.\(tableName).find({}).limit(\(pageSize))"
568+
let escaped = tableName.replacingOccurrences(of: "\\", with: "\\\\").replacingOccurrences(of: "\"", with: "\\\"")
569+
query = "db[\"\(escaped)\"].find({}).limit(\(pageSize))"
569570
} else {
570571
let quotedName = databaseType.quoteIdentifier(tableName)
571572
query = "SELECT * FROM \(quotedName) LIMIT \(pageSize);"
@@ -599,7 +600,8 @@ final class QueryTabManager: ObservableObject {
599600
let pageSize = AppSettingsManager.shared.dataGrid.defaultPageSize
600601
let query: String
601602
if databaseType == .mongodb {
602-
query = "db.\(tableName).find({}).limit(\(pageSize))"
603+
let escaped = tableName.replacingOccurrences(of: "\\", with: "\\\\").replacingOccurrences(of: "\"", with: "\\\"")
604+
query = "db[\"\(escaped)\"].find({}).limit(\(pageSize))"
603605
} else {
604606
let quotedName = databaseType.quoteIdentifier(tableName)
605607
query = "SELECT * FROM \(quotedName) LIMIT \(pageSize);"

TablePro/ViewModels/SidebarViewModel.swift

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
//
88

99
import Combine
10-
import os
1110
import SwiftUI
1211

1312
// MARK: - TableFetcher Protocol
@@ -34,7 +33,9 @@ struct LiveTableFetcher: TableFetcher {
3433
return cached
3534
}
3635
}
37-
guard let driver = await DatabaseManager.shared.driver(for: connectionId) else { return [] }
36+
guard let driver = await DatabaseManager.shared.driver(for: connectionId) else {
37+
return []
38+
}
3839
return try await driver.fetchTables()
3940
}
4041
}
@@ -43,8 +44,6 @@ struct LiveTableFetcher: TableFetcher {
4344

4445
@MainActor
4546
final class SidebarViewModel: ObservableObject {
46-
private static let logger = Logger(subsystem: "com.TablePro", category: "SidebarViewModel")
47-
4847
// MARK: - Published State
4948

5049
@Published var isLoading = false
@@ -82,6 +81,7 @@ final class SidebarViewModel: ObservableObject {
8281
private let tableFetcher: TableFetcher
8382
private var cancellables = Set<AnyCancellable>()
8483
private var hasSetupNotifications = false
84+
private var loadTask: Task<Void, Never>?
8585

8686
// MARK: - Convenience Accessors
8787

@@ -167,7 +167,7 @@ final class SidebarViewModel: ObservableObject {
167167
.receive(on: DispatchQueue.main)
168168
.sink { [weak self] _ in
169169
Task { @MainActor in
170-
self?.loadTables()
170+
self?.forceLoadTables()
171171
}
172172
}
173173
.store(in: &cancellables)
@@ -201,11 +201,18 @@ final class SidebarViewModel: ObservableObject {
201201
guard !isLoading else { return }
202202
isLoading = true
203203
errorMessage = nil
204-
Task {
204+
loadTask = Task {
205205
await loadTablesAsync()
206206
}
207207
}
208208

209+
func forceLoadTables() {
210+
loadTask?.cancel()
211+
loadTask = nil
212+
isLoading = false
213+
loadTables()
214+
}
215+
209216
private func loadTablesAsync() async {
210217
let previousSelectedName = selectedTables.first?.name
211218

TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,16 @@ extension MainContentCoordinator {
272272
runQuery()
273273
}
274274
} else if connection.type == .mongodb {
275-
// MongoDB: just update the database name — driver reads it for every operation
275+
// MongoDB: update the driver's connection so fetchTables/execute use the new database
276+
if let mongoDriver = driver as? MongoDBDriver {
277+
mongoDriver.switchDatabase(to: database)
278+
}
279+
280+
// Also update metadata driver if present
281+
if let metaDriver = DatabaseManager.shared.metadataDriver(for: connectionId) as? MongoDBDriver {
282+
metaDriver.switchDatabase(to: database)
283+
}
284+
276285
DatabaseManager.shared.updateSession(connectionId) { session in
277286
var updatedConnection = session.connection
278287
updatedConnection.database = database
@@ -295,6 +304,8 @@ extension MainContentCoordinator {
295304

296305
await loadSchema()
297306

307+
NotificationCenter.default.post(name: .refreshData, object: nil)
308+
298309
if let currentTab = tabManager.selectedTab, currentTab.tabType == .table {
299310
runQuery()
300311
}

TablePro/Views/Main/Extensions/MainContentCoordinator+TableOperations.swift

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,8 @@ extension MainContentCoordinator {
126126
"DELETE FROM sqlite_sequence WHERE name = '\(escapedName)'"
127127
]
128128
case .mongodb:
129-
return ["db.\(tableName).deleteMany({})"]
129+
let escaped = tableName.replacingOccurrences(of: "\\", with: "\\\\").replacingOccurrences(of: "\"", with: "\\\"")
130+
return ["db[\"\(escaped)\"].deleteMany({})"]
130131
}
131132
}
132133

@@ -139,7 +140,8 @@ extension MainContentCoordinator {
139140
case .mysql, .mariadb, .sqlite:
140141
return "DROP \(keyword) \(quotedName)"
141142
case .mongodb:
142-
return "db.\(tableName).drop()"
143+
let escaped = tableName.replacingOccurrences(of: "\\", with: "\\\\").replacingOccurrences(of: "\"", with: "\\\"")
144+
return "db[\"\(escaped)\"].drop()"
143145
}
144146
}
145147
}

TablePro/Views/Sidebar/SidebarView.swift

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@ struct SidebarView: View {
7373
}
7474
.frame(minWidth: 280)
7575
.onChange(of: tables) { _, newTables in
76-
if newTables.isEmpty && DatabaseManager.shared.activeSessions[connectionId] != nil && !viewModel.isLoading {
76+
let hasSession = DatabaseManager.shared.activeSessions[connectionId] != nil
77+
if newTables.isEmpty && hasSession && !viewModel.isLoading {
7778
viewModel.loadTables()
7879
}
7980
}

0 commit comments

Comments
 (0)