|
| 1 | +// |
| 2 | +// RedisKeyTreeViewModel.swift |
| 3 | +// TablePro |
| 4 | +// |
| 5 | + |
| 6 | +import Foundation |
| 7 | +import Observation |
| 8 | +import os |
| 9 | + |
| 10 | +@MainActor @Observable |
| 11 | +internal final class RedisKeyTreeViewModel { |
| 12 | + private static let logger = Logger(subsystem: "com.TablePro", category: "RedisKeyTree") |
| 13 | + private static let maxKeys = 50_000 |
| 14 | + |
| 15 | + var rootNodes: [RedisKeyNode] = [] |
| 16 | + var expandedPrefixes: Set<String> = [] |
| 17 | + var isLoading = false |
| 18 | + var isTruncated = false |
| 19 | + var separator: String = ":" |
| 20 | + |
| 21 | + private(set) var allKeys: [(key: String, type: String)] = [] |
| 22 | + |
| 23 | + /// Test-only setter for allKeys |
| 24 | + var allKeysForTesting: [(key: String, type: String)] { |
| 25 | + get { allKeys } |
| 26 | + set { allKeys = newValue } |
| 27 | + } |
| 28 | + |
| 29 | + func loadKeys(connectionId: UUID, database: String, separator: String) async { |
| 30 | + self.separator = separator |
| 31 | + isLoading = true |
| 32 | + isTruncated = false |
| 33 | + defer { isLoading = false } |
| 34 | + |
| 35 | + guard let driver = DatabaseManager.shared.driver(for: connectionId) else { |
| 36 | + clear() |
| 37 | + return |
| 38 | + } |
| 39 | + |
| 40 | + do { |
| 41 | + // Use KEYS command for simplicity — returns all keys matching pattern |
| 42 | + let result = try await driver.execute(query: "KEYS *") |
| 43 | + |
| 44 | + let keyColumnIndex = result.columns.firstIndex(of: "Key") ?? 0 |
| 45 | + let typeColumnIndex = result.columns.firstIndex(of: "Type") ?? 1 |
| 46 | + |
| 47 | + var keys: [(key: String, type: String)] = [] |
| 48 | + for row in result.rows { |
| 49 | + guard keyColumnIndex < row.count, |
| 50 | + let keyName = row[keyColumnIndex] else { continue } |
| 51 | + let keyType = typeColumnIndex < row.count ? (row[typeColumnIndex] ?? "string") : "string" |
| 52 | + keys.append((key: keyName, type: keyType)) |
| 53 | + if keys.count >= Self.maxKeys { break } |
| 54 | + } |
| 55 | + |
| 56 | + isTruncated = keys.count >= Self.maxKeys |
| 57 | + allKeys = keys |
| 58 | + rootNodes = Self.buildTree(keys: keys, separator: separator) |
| 59 | + } catch { |
| 60 | + Self.logger.error("Failed to load Redis keys: \(error.localizedDescription, privacy: .public)") |
| 61 | + clear() |
| 62 | + } |
| 63 | + } |
| 64 | + |
| 65 | + func clear() { |
| 66 | + rootNodes = [] |
| 67 | + allKeys = [] |
| 68 | + expandedPrefixes = [] |
| 69 | + isTruncated = false |
| 70 | + } |
| 71 | + |
| 72 | + func displayNodes(searchText: String) -> [RedisKeyNode] { |
| 73 | + guard !searchText.isEmpty else { return rootNodes } |
| 74 | + |
| 75 | + let filtered = allKeys.filter { $0.key.localizedCaseInsensitiveContains(searchText) } |
| 76 | + if filtered.isEmpty { return [] } |
| 77 | + |
| 78 | + return Self.buildTree(keys: filtered, separator: separator) |
| 79 | + } |
| 80 | + |
| 81 | + // MARK: - Tree Building (Pure Function) |
| 82 | + |
| 83 | + static func buildTree(keys: [(key: String, type: String)], separator: String) -> [RedisKeyNode] { |
| 84 | + guard !separator.isEmpty else { |
| 85 | + return keys.sorted { $0.key < $1.key } |
| 86 | + .map { .key(name: $0.key, fullKey: $0.key, keyType: $0.type) } |
| 87 | + } |
| 88 | + |
| 89 | + var root = TrieNode() |
| 90 | + for entry in keys { |
| 91 | + let parts = entry.key.components(separatedBy: separator) |
| 92 | + root.insert(parts: parts, fullKey: entry.key, keyType: entry.type) |
| 93 | + } |
| 94 | + |
| 95 | + return root.toRedisKeyNodes(parentPrefix: "", separator: separator) |
| 96 | + } |
| 97 | +} |
| 98 | + |
| 99 | +// MARK: - Trie for Tree Building |
| 100 | + |
| 101 | +private class TrieNode { |
| 102 | + var children: [String: TrieNode] = [:] |
| 103 | + var leafKeys: [(fullKey: String, keyType: String)] = [] |
| 104 | + |
| 105 | + func insert(parts: [String], fullKey: String, keyType: String) { |
| 106 | + guard !parts.isEmpty else { |
| 107 | + leafKeys.append((fullKey: fullKey, keyType: keyType)) |
| 108 | + return |
| 109 | + } |
| 110 | + |
| 111 | + if parts.count == 1 { |
| 112 | + leafKeys.append((fullKey: fullKey, keyType: keyType)) |
| 113 | + } else { |
| 114 | + let segment = parts[0] |
| 115 | + let child = children[segment] ?? TrieNode() |
| 116 | + children[segment] = child |
| 117 | + child.insert(parts: Array(parts.dropFirst()), fullKey: fullKey, keyType: keyType) |
| 118 | + } |
| 119 | + } |
| 120 | + |
| 121 | + func toRedisKeyNodes(parentPrefix: String, separator: String) -> [RedisKeyNode] { |
| 122 | + var nodes: [RedisKeyNode] = [] |
| 123 | + |
| 124 | + let sortedChildren = children.sorted { $0.key < $1.key } |
| 125 | + for (segment, child) in sortedChildren { |
| 126 | + let fullPrefix = parentPrefix.isEmpty ? "\(segment)\(separator)" : "\(parentPrefix)\(segment)\(separator)" |
| 127 | + let childNodes = child.toRedisKeyNodes(parentPrefix: fullPrefix, separator: separator) |
| 128 | + let keyCount = child.countLeafKeys() |
| 129 | + |
| 130 | + if !childNodes.isEmpty || !child.leafKeys.isEmpty { |
| 131 | + nodes.append(.namespace( |
| 132 | + name: segment, |
| 133 | + fullPrefix: fullPrefix, |
| 134 | + children: childNodes, |
| 135 | + keyCount: keyCount |
| 136 | + )) |
| 137 | + } |
| 138 | + } |
| 139 | + |
| 140 | + let sortedLeafs = leafKeys.sorted { $0.fullKey < $1.fullKey } |
| 141 | + for leaf in sortedLeafs { |
| 142 | + let displayName: String |
| 143 | + if parentPrefix.isEmpty { |
| 144 | + displayName = leaf.fullKey |
| 145 | + } else { |
| 146 | + displayName = String(leaf.fullKey.dropFirst(parentPrefix.count)) |
| 147 | + } |
| 148 | + nodes.append(.key(name: displayName, fullKey: leaf.fullKey, keyType: leaf.keyType)) |
| 149 | + } |
| 150 | + |
| 151 | + return nodes |
| 152 | + } |
| 153 | + |
| 154 | + func countLeafKeys() -> Int { |
| 155 | + var count = leafKeys.count |
| 156 | + for child in children.values { |
| 157 | + count += child.countLeafKeys() |
| 158 | + } |
| 159 | + return count |
| 160 | + } |
| 161 | +} |
0 commit comments