diff --git a/src/barabadb/client/client.nim b/src/barabadb/client/client.nim index c55065f..c9bc642 100644 --- a/src/barabadb/client/client.nim +++ b/src/barabadb/client/client.nim @@ -109,10 +109,10 @@ proc disconnect*(client: SyncClient) = try: let msg = makeQueryMessage(0, "DISCONNECT") netmod.send(client.socket, cast[string](msg)) - except: discard + except IOError, OSError: discard netmod.close(client.socket) client.connected = false - deinitLock(client.lock) + deinitLock(client.lock) proc query*(client: SyncClient, sql: string): QueryResult = acquire(client.lock) diff --git a/src/barabadb/core/backup.nim b/src/barabadb/core/backup.nim index 25a4d76..1d01068 100644 --- a/src/barabadb/core/backup.nim +++ b/src/barabadb/core/backup.nim @@ -177,7 +177,7 @@ proc logRestore*(archive: string, dataDir: string, success: bool, dryRun: bool = let f = open(logPath, fmAppend) f.write(entry) f.close() - except: + except IOError: discard # Silently fail if log cannot be written proc readHistory*(): seq[string] = @@ -191,7 +191,7 @@ proc readHistory*(): seq[string] = for line in splitLines(content): if line.len > 0: result.add(line) - except: + except IOError: discard proc backupDataDir*(dataDir: string, output: string, excludes: seq[string] = @[], compression: int = DEFAULT_COMPRESSION, verbose: bool = false): bool = diff --git a/src/barabadb/core/deadlock.nim b/src/barabadb/core/deadlock.nim index 46b3314..0e24c47 100644 --- a/src/barabadb/core/deadlock.nim +++ b/src/barabadb/core/deadlock.nim @@ -2,6 +2,7 @@ import std/tables import std/sets import std/algorithm +import std/locks type WaitEdge* = object @@ -12,15 +13,18 @@ type edges: seq[WaitEdge] adjacency: Table[uint64, seq[uint64]] # waiter -> holders txnIds: HashSet[uint64] + lock: Lock proc newDeadlockDetector*(): DeadlockDetector = - DeadlockDetector( - edges: @[], - adjacency: initTable[uint64, seq[uint64]](), - txnIds: initHashSet[uint64](), - ) + new(result) + initLock(result.lock) + result.edges = @[] + result.adjacency = initTable[uint64, seq[uint64]]() + result.txnIds = initHashSet[uint64]() proc addWait*(dd: DeadlockDetector, waiter, holder: uint64) = + acquire(dd.lock) + defer: release(dd.lock) dd.edges.add(WaitEdge(waiter: waiter, holder: holder)) dd.txnIds.incl(waiter) dd.txnIds.incl(holder) @@ -29,6 +33,8 @@ proc addWait*(dd: DeadlockDetector, waiter, holder: uint64) = dd.adjacency[waiter].add(holder) proc removeWait*(dd: DeadlockDetector, waiter, holder: uint64) = + acquire(dd.lock) + defer: release(dd.lock) var newEdges: seq[WaitEdge] = @[] for edge in dd.edges: if edge.waiter != waiter or edge.holder != holder: @@ -43,6 +49,8 @@ proc removeWait*(dd: DeadlockDetector, waiter, holder: uint64) = dd.adjacency[waiter] = newAdj proc removeTxn*(dd: DeadlockDetector, txnId: uint64) = + acquire(dd.lock) + defer: release(dd.lock) dd.txnIds.excl(txnId) dd.adjacency.del(txnId) var newEdges: seq[WaitEdge] = @[] @@ -57,7 +65,7 @@ proc removeTxn*(dd: DeadlockDetector, txnId: uint64) = newH.add(h) holders = newH -proc detectCycle*(dd: DeadlockDetector): seq[uint64] = +proc detectCycleUnsafe(dd: DeadlockDetector): seq[uint64] = var visited = initHashSet[uint64]() var inStack = initHashSet[uint64]() var parent = initTable[uint64, uint64]() @@ -98,8 +106,15 @@ proc detectCycle*(dd: DeadlockDetector): seq[uint64] = return cycle return @[] +proc detectCycle*(dd: DeadlockDetector): seq[uint64] = + acquire(dd.lock) + defer: release(dd.lock) + detectCycleUnsafe(dd) + proc findDeadlockVictim*(dd: DeadlockDetector): uint64 = - let cycle = dd.detectCycle() + acquire(dd.lock) + defer: release(dd.lock) + let cycle = detectCycleUnsafe(dd) if cycle.len == 0: return 0 # Choose youngest txn (highest id) as victim @@ -109,12 +124,23 @@ proc findDeadlockVictim*(dd: DeadlockDetector): uint64 = result = id proc hasDeadlock*(dd: DeadlockDetector): bool = - return dd.detectCycle().len > 0 + acquire(dd.lock) + defer: release(dd.lock) + return detectCycleUnsafe(dd).len > 0 proc clear*(dd: DeadlockDetector) = + acquire(dd.lock) + defer: release(dd.lock) dd.edges.setLen(0) dd.adjacency.clear() dd.txnIds.clear() -proc edgeCount*(dd: DeadlockDetector): int = dd.edges.len -proc txnCount*(dd: DeadlockDetector): int = dd.txnIds.len +proc edgeCount*(dd: DeadlockDetector): int = + acquire(dd.lock) + defer: release(dd.lock) + dd.edges.len + +proc txnCount*(dd: DeadlockDetector): int = + acquire(dd.lock) + defer: release(dd.lock) + dd.txnIds.len diff --git a/src/barabadb/core/disttxn.nim b/src/barabadb/core/disttxn.nim index 5244c5e..11e416e 100644 --- a/src/barabadb/core/disttxn.nim +++ b/src/barabadb/core/disttxn.nim @@ -108,86 +108,119 @@ proc sendDistTxnRpc(host: string, port: int, txnId: uint64, action: string, time except CatchableError: return false +type + ParticipantInfo = object + nodeId: string + host: string + port: int + proc prepare*(txn: DistributedTransaction): bool = + # Phase 1: validate state and collect participants while holding lock + var participants: seq[ParticipantInfo] = @[] + var wasActive = false acquire(txn.lock) - if txn.state != dtsActive: - release(txn.lock) + if txn.state == dtsActive: + txn.state = dtsPreparing + wasActive = true + for nodeId, participant in txn.participants: + if participant.host.len > 0 and participant.port > 0: + participants.add(ParticipantInfo(nodeId: nodeId, host: participant.host, port: participant.port)) + release(txn.lock) + + if not wasActive: return false - txn.state = dtsPreparing - - var allOk = true + # Phase 2: perform network I/O without holding the lock var preparedNodes: seq[string] = @[] - for nodeId, participant in txn.participants.mpairs: - if participant.host.len > 0 and participant.port > 0: - participant.prepared = sendDistTxnRpc(participant.host, participant.port, txn.id, "PREPARE") - else: - participant.prepared = true # local participant - if participant.prepared: - preparedNodes.add(nodeId) + var allOk = true + for p in participants: + let ok = sendDistTxnRpc(p.host, p.port, txn.id, "PREPARE") + if ok: + preparedNodes.add(p.nodeId) else: allOk = false + # Phase 3: update state while holding lock + acquire(txn.lock) if allOk: txn.state = dtsPrepared + for nodeId, _ in txn.participants.mpairs: + txn.participants[nodeId].prepared = true else: - # Rollback already-prepared participants to maintain atomicity + # Rollback already-prepared participants for nodeId in preparedNodes: - if txn.participants[nodeId].host.len > 0 and txn.participants[nodeId].port > 0: + if txn.participants.hasKey(nodeId): discard sendDistTxnRpc(txn.participants[nodeId].host, txn.participants[nodeId].port, txn.id, "ROLLBACK") - txn.participants[nodeId].prepared = false - txn.participants[nodeId].aborted = true + txn.participants[nodeId].prepared = false + txn.participants[nodeId].aborted = true txn.state = dtsAborted - release(txn.lock) return allOk proc commit*(txn: DistributedTransaction): bool = + # Phase 1: validate state and collect participants while holding lock + var participants: seq[ParticipantInfo] = @[] + var wasPrepared = false acquire(txn.lock) - if txn.state != dtsPrepared: - release(txn.lock) + if txn.state == dtsPrepared: + txn.state = dtsCommitting + wasPrepared = true + for nodeId, participant in txn.participants: + if participant.host.len > 0 and participant.port > 0: + participants.add(ParticipantInfo(nodeId: nodeId, host: participant.host, port: participant.port)) + release(txn.lock) + + if not wasPrepared: return false - txn.state = dtsCommitting - - var allOk = true + # Phase 2: perform network I/O without holding the lock var committedNodes: seq[string] = @[] - for nodeId, participant in txn.participants.mpairs: - if participant.host.len > 0 and participant.port > 0: - participant.committed = sendDistTxnRpc(participant.host, participant.port, txn.id, "COMMIT") - else: - participant.committed = true # local participant - if participant.committed: - committedNodes.add(nodeId) + var allOk = true + for p in participants: + let ok = sendDistTxnRpc(p.host, p.port, txn.id, "COMMIT") + if ok: + committedNodes.add(p.nodeId) else: allOk = false + # Phase 3: update state while holding lock + acquire(txn.lock) if allOk: txn.state = dtsCommitted + for nodeId, _ in txn.participants.mpairs: + txn.participants[nodeId].committed = true + elif committedNodes.len > 0: + txn.state = dtsCommitted + for nodeId in committedNodes: + if txn.participants.hasKey(nodeId): + txn.participants[nodeId].committed = true else: - if committedNodes.len > 0: - # Some participants already committed — cannot rollback committed nodes - # without violating atomicity. Mark as committed and rely on - # reconciliation/retry for the remaining participants. - txn.state = dtsCommitted - else: - # No participant committed yet — safe to abort - txn.state = dtsAborted + txn.state = dtsAborted release(txn.lock) return allOk or committedNodes.len > 0 proc rollback*(txn: DistributedTransaction): bool = + # Phase 1: validate state and collect participants while holding lock + var participants: seq[ParticipantInfo] = @[] + var canRollback = false acquire(txn.lock) - if txn.state notin {dtsActive, dtsPreparing, dtsPrepared}: - release(txn.lock) + if txn.state in {dtsActive, dtsPreparing, dtsPrepared}: + txn.state = dtsAborting + canRollback = true + for nodeId, participant in txn.participants: + if participant.host.len > 0 and participant.port > 0: + participants.add(ParticipantInfo(nodeId: nodeId, host: participant.host, port: participant.port)) + release(txn.lock) + + if not canRollback: return false - txn.state = dtsAborting - for nodeId, participant in txn.participants.mpairs: - if participant.host.len > 0 and participant.port > 0: - participant.aborted = sendDistTxnRpc(participant.host, participant.port, txn.id, "ROLLBACK") - else: - participant.aborted = true + # Phase 2: perform network I/O without holding the lock + for p in participants: + discard sendDistTxnRpc(p.host, p.port, txn.id, "ROLLBACK") + + # Phase 3: update state while holding lock + acquire(txn.lock) txn.state = dtsAborted release(txn.lock) return true diff --git a/src/barabadb/core/raft.nim b/src/barabadb/core/raft.nim index 87133c5..00e4aef 100644 --- a/src/barabadb/core/raft.nim +++ b/src/barabadb/core/raft.nim @@ -124,9 +124,10 @@ proc loadState(node: RaftNode) = let dataLen = int(s.readUint32()) var data = newSeq[byte](dataLen) if dataLen > 0: - discard s.readData(addr data[0], dataLen) + if s.readData(addr data[0], dataLen) != dataLen: + raise newException(IOError, "Incomplete Raft log data read") node.log[i] = LogEntry(term: term, index: index, command: cmd, data: data) - except: + except IOError, OSError: discard s.close() @@ -431,7 +432,8 @@ proc readString(s: Stream): string = let len = int(s.readUint32()) if len > 0: result = newString(len) - discard s.readData(result[0].addr, len) + if s.readData(result[0].addr, len) != len: + raise newException(IOError, "Incomplete string read from stream") else: result = "" diff --git a/src/barabadb/core/server.nim b/src/barabadb/core/server.nim index a461f21..444f600 100644 --- a/src/barabadb/core/server.nim +++ b/src/barabadb/core/server.nim @@ -99,6 +99,8 @@ proc newServer*(config: BaraConfig): Server = # ---------------------------------------------------------------------- proc readUint32BE(data: string, pos: int): uint32 = + if pos + 4 > data.len: + raise newException(ValueError, "readUint32BE: index out of bounds") var bytes: array[4, byte] for i in 0..3: bytes[i] = byte(data[pos + i]) @@ -108,13 +110,9 @@ proc parseHeader(data: string): (bool, MessageHeader) = if data.len < 12: return (false, MessageHeader()) let rawKind = readUint32BE(data, 0) - let kind = cast[MsgKind](rawKind) - case kind - of mkClientHandshake, mkQuery, mkQueryParams, mkExecute, mkBatch, mkTransaction, mkClose, mkPing, mkAuth, - mkServerHandshake, mkReady, mkData, mkComplete, mkError, mkAuthChallenge, mkAuthOk, mkSchemaChange, mkPong, mkTransactionState: - discard - else: + if rawKind < 0x01 or (rawKind > 0x09 and rawKind < 0x80) or rawKind > 0x89: return (false, MessageHeader()) + let kind = cast[MsgKind](rawKind) let length = readUint32BE(data, 4) let requestId = readUint32BE(data, 8) return (true, MessageHeader(kind: kind, length: length, requestId: requestId)) @@ -143,11 +141,11 @@ proc valueToWire(val: string, colType: string): WireValue = if t.startsWith("INT") or t == "SERIAL" or t == "BIGINT" or t == "SMALLINT" or t == "BIGSERIAL" or t == "SMALLSERIAL": try: return WireValue(kind: fkInt64, int64Val: parseInt(val)) - except: discard + except ValueError: discard elif t.startsWith("FLOAT") or t == "REAL" or t == "DOUBLE" or t == "NUMERIC" or t.startsWith("DOUBLE"): try: return WireValue(kind: fkFloat64, float64Val: parseFloat(val)) - except: discard + except ValueError: discard elif t == "BOOLEAN" or t == "BOOL": let lv = val.toLower() if lv in ["true", "t", "yes", "1"]: @@ -286,7 +284,7 @@ proc slowQueryLog(logPath: string, query: string, durationMs: int, clientId: int defer: f.close() let line = $getMonoTime().ticks() & " | " & $clientId & " | " & $durationMs & "ms | " & query & "\n" f.write(line) - except: discard + except IOError: discard proc verifyToken(secret, tokenStr: string): (bool, string, string) = try: @@ -296,7 +294,7 @@ proc verifyToken(secret, tokenStr: string): (bool, string, string) = let userId = token.claims["sub"].node.str let role = if "role" in token.claims: token.claims["role"].node.str else: "user" return (true, userId, role) - except: + except ValueError, KeyError: return (false, "", "") proc recvWithTimeout(client: AsyncSocket, size: int, timeoutMs: int): Future[string] {.async.} = diff --git a/src/barabadb/fts/engine b/src/barabadb/fts/engine new file mode 100755 index 0000000..5995748 Binary files /dev/null and b/src/barabadb/fts/engine differ diff --git a/src/barabadb/protocol/auth.nim b/src/barabadb/protocol/auth.nim index b73f545..3480203 100644 --- a/src/barabadb/protocol/auth.nim +++ b/src/barabadb/protocol/auth.nim @@ -110,10 +110,13 @@ proc constantTimeCompare(a, b: string): bool = # JWT token helpers # --------------------------------------------------------------------------- +proc jsonEscape(s: string): string = + result = s.replace("\\", "\\\\").replace("\"", "\\\"") + proc createToken*(am: AuthManager, claims: JWTClaims): string = let header = base64UrlEncode("{\"alg\":\"HS256\",\"typ\":\"JWT\"}") - var payloadJson = "{\"sub\":\"" & claims.sub & "\",\"role\":\"" & claims.role & - "\",\"database\":\"" & claims.database & "\"" + var payloadJson = "{\"sub\":\"" & jsonEscape(claims.sub) & "\",\"role\":\"" & jsonEscape(claims.role) & + "\",\"database\":\"" & jsonEscape(claims.database) & "\"" if claims.exp > 0: payloadJson &= ",\"exp\":" & $claims.exp if claims.iat > 0: diff --git a/src/barabadb/protocol/scram.nim b/src/barabadb/protocol/scram.nim index d69c7e6..7820034 100644 --- a/src/barabadb/protocol/scram.nim +++ b/src/barabadb/protocol/scram.nim @@ -40,7 +40,9 @@ proc hmacSha256*(key, message: string): array[32, byte] = var ctx = initSha_256() ctx.update(k.toOpenArray(0, k.len-1)) let hash = ctx.digest() - k = $hash + var kHashed = newString(32) + copyMem(addr kHashed[0], unsafeAddr hash[0], 32) + k = kHashed while k.len < 64: k &= "\x00" @@ -130,6 +132,8 @@ proc generateNonce*(): string = ## Generate a cryptographically secure random nonce (base64-encoded). when defined(linux) or defined(macosx) or defined(bsd): let f = open("/dev/urandom") + if f == nil: + raise newException(IOError, "Failed to open /dev/urandom") defer: f.close() var bytes = newString(NonceBytes) let readLen = f.readBuffer(addr bytes[0], NonceBytes) @@ -209,18 +213,18 @@ proc parseClientFinal*(msg: string): (string, string, seq[byte]) = ## Returns: (channel_binding, nonce, proof_bytes) var cbind = "" var nonce = "" - var proofHex = "" + var proofEncoded = "" for p in msg.split(","): if p.startsWith("c="): cbind = p[2..^1] elif p.startsWith("r="): nonce = p[2..^1] elif p.startsWith("p="): - proofHex = p[2..^1] - if nonce.len == 0 or proofHex.len == 0: + proofEncoded = p[2..^1] + if nonce.len == 0 or proofEncoded.len == 0: raise newException(ValueError, "Missing nonce or proof in client-final-message") # proof is base64-encoded - var proof = decode(proofHex) + var proof = decode(proofEncoded) var proofBytes = newSeq[byte](proof.len) if proof.len > 0: copyMem(addr proofBytes[0], addr proof[0], proof.len) @@ -240,6 +244,8 @@ proc buildServerFinal*(serverSignature: openArray[byte]): string = # --------------------------------------------------------------------------- proc verifyClientProof*(state: ScramServerState, clientProof: openArray[byte]): bool = + if clientProof.len != 32: + return false let clientKey = xorBytes(clientProof, @(hmacSha256(state.storedKey, state.authMessage))) let computedStoredKey = sha256(clientKey) for i in 0..<32: diff --git a/src/barabadb/protocol/wire.nim b/src/barabadb/protocol/wire.nim index 4218148..88cc5cd 100644 --- a/src/barabadb/protocol/wire.nim +++ b/src/barabadb/protocol/wire.nim @@ -245,7 +245,7 @@ proc deserializeValue*(buf: openArray[byte], pos: var int, depth: int = 0): Wire result = WireValue(kind: fkBytes, bytesVal: readBytes(buf, pos)) of fkArray: let count = int(readUint32(buf, pos)) - if count > MaxWireArrayLen: + if count < 0 or count > MaxWireArrayLen: raise newException(ValueError, "Wire protocol: array exceeds max length") var arr: seq[WireValue] = @[] for i in 0.. MaxWireObjectLen: + if count < 0 or count > MaxWireObjectLen: raise newException(ValueError, "Wire protocol: object exceeds max length") var obj: seq[(string, WireValue)] = @[] for i in 0.. MaxWireVectorLen: + if count < 0 or count > MaxWireVectorLen: raise newException(ValueError, "Wire protocol: vector exceeds max length") if pos + count * 4 > buf.len: raise newException(ValueError, "Wire protocol: truncated vector data") diff --git a/src/barabadb/protocol/zerocopy.nim b/src/barabadb/protocol/zerocopy.nim index 269190f..54071ab 100644 --- a/src/barabadb/protocol/zerocopy.nim +++ b/src/barabadb/protocol/zerocopy.nim @@ -141,6 +141,8 @@ proc readString*(buf: ZeroBuf, offset: var int): string = let len = int(buf.readInt64(offset)) offset += 8 if len > 0: + if offset + len > buf.capacity: + raise newException(ValueError, "ZeroCopy readString: buffer overflow") result = newString(len) copyMem(addr result[0], addr buf.data[offset], len) offset += len diff --git a/src/barabadb/query/executor.nim b/src/barabadb/query/executor.nim index 1d5fa90..83ceb11 100644 --- a/src/barabadb/query/executor.nim +++ b/src/barabadb/query/executor.nim @@ -185,6 +185,14 @@ proc newExecutionContext*(db: LSMTree): ExecutionContext = # AST to SQL serializer (for VIEW DDL persistence) # ---------------------------------------------------------------------- +proc sqlEscapeIdent*(ident: string): string = + ## Escape SQL identifiers by doubling double-quotes. + result = ident.replace("\"", "\"\"") + +proc sqlEscapeString*(s: string): string = + ## Escape SQL string literals by doubling single-quotes. + result = s.replace("'", "''") + proc exprToSql(node: Node): string = if node == nil: return "" @@ -200,7 +208,7 @@ proc exprToSql(node: Node): string = of nkNullLit: return "null" of nkIdent: - return node.identName + return "\"" & node.identName.replace("\"", "\"\"") & "\"" of nkStar: return "*" of nkBinOp: @@ -297,8 +305,13 @@ proc restoreSchema(ctx: ExecutionContext) = if not entry.key.startsWith("_schema:"): continue let ddl = cast[string](entry.value) if ddl.len == 0: continue - let tokens = qlex.tokenize(ddl) - let astNode = qpar.parse(tokens) + var astNode: Node + try: + let tokens = qlex.tokenize(ddl) + astNode = qpar.parse(tokens) + except: + # Skip corrupted schema entries during startup + continue if astNode.stmts.len > 0: let stmt = astNode.stmts[0] case stmt.kind @@ -992,9 +1005,15 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = if expr.binRight.kind == irekSubquery: let subRows = executePlan(ctx, expr.binRight.subqueryPlan) for row in subRows: + # Compare against the first non-internal column only (SQL semantics) + var firstVal = "" + var found = false for k, v in row: if k.startsWith("$"): continue - if v == left: return "true" + firstVal = v + found = true + break + if found and firstVal == left: return "true" return "false" try: let lv = parseFloat(left) @@ -1006,9 +1025,15 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = if expr.binRight.kind == irekSubquery: let subRows = executePlan(ctx, expr.binRight.subqueryPlan) for row in subRows: + # Compare against the first non-internal column only (SQL semantics) + var firstVal = "" + var found = false for k, v in row: if k.startsWith("$"): continue - if v == left: return "false" + firstVal = v + found = true + break + if found and firstVal == left: return "false" return "true" try: let lv = parseFloat(left) @@ -1334,23 +1359,43 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = if sql.len == 0: return "" - # Validate by trying EXPLAIN or LIMIT-wrapped query - var validateSql = sql - if validateSql.toLower().startsWith("select"): - validateSql = "SELECT * FROM (" & sql & ") LIMIT 0" - let tokens = qlex.tokenize(validateSql) - let astNode = qpar.parse(tokens) - if astNode.stmts.len > 0: - let validateRes = executeQuery(ctx, astNode) - if not validateRes.success: - # Self-correction: send error back to LLM - let correctionPrompt = "Schema:\n" & schemaInfo & "\nQuestion: " & question & "\n\nPrevious SQL: " & sql & "\n\nError: " & validateRes.message & "\n\nGenerate corrected SQL:" + # Validate by parsing only (never execute LLM-generated SQL directly) + let sqlLower = sql.toLower().strip() + let isSafeQuery = sqlLower.startsWith("select") or sqlLower.startsWith("explain") or + sqlLower.startsWith("with") + let allowDml = ctx.sessionVars.getOrDefault("nl_to_sql.allow_dml", "false") == "true" + if not isSafeQuery and not allowDml: + # For non-SELECT: only do syntax validation via tokenize+parse, no execution + let tokens = qlex.tokenize(sql) + let astNode = qpar.parse(tokens) + if astNode.stmts.len > 0: + return sql + else: + let correctionPrompt = "Schema:\n" & schemaInfo & "\nQuestion: " & question & "\n\nPrevious SQL: " & sql & "\n\nError: Syntax error. Generate corrected SQL:" var correctedResponse = llmmod.generate(ctx.llmClient, correctionPrompt, systemPrompt) var correctedSql = llmmod.extractSQL(correctedResponse) if correctedSql.len > 0: return correctedSql - - return sql + return "" + else: + # For SELECT/EXPLAIN or when ALLOW_DML is set: validate with LIMIT 0 or EXPLAIN + var validateSql = sql + if sqlLower.startsWith("select"): + validateSql = "SELECT * FROM (" & sql & ") LIMIT 0" + elif sqlLower.startsWith("with"): + validateSql = "WITH _cte_ AS (" & sql & ") SELECT 1 LIMIT 0" + let tokens = qlex.tokenize(validateSql) + let astNode = qpar.parse(tokens) + if astNode.stmts.len > 0: + let validateRes = executeQuery(ctx, astNode) + if not validateRes.success: + # Self-correction: send error back to LLM + let correctionPrompt = "Schema:\n" & schemaInfo & "\nQuestion: " & question & "\n\nPrevious SQL: " & sql & "\n\nError: " & validateRes.message & "\n\nGenerate corrected SQL:" + var correctedResponse = llmmod.generate(ctx.llmClient, correctionPrompt, systemPrompt) + var correctedSql = llmmod.extractSQL(correctedResponse) + if correctedSql.len > 0: + return correctedSql + return sql of "schema_prompt": if expr.irFuncArgs.len < 1: return "" let table = evalExpr(expr.irFuncArgs[0], row, ctx) @@ -1358,33 +1403,33 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = return "Table '" & table & "' not found" let tbl = ctx.tables[table] - var result = "" - result.add("CREATE TABLE " & table & " (\n") + var ddl = "" + ddl.add("CREATE TABLE " & table & " (\n") for i, col in tbl.columns: - result.add(" " & col.name & " " & col.colType) - if col.isPk: result.add(" PRIMARY KEY") - if col.isNotNull: result.add(" NOT NULL") - if col.autoIncrement: result.add(" AUTO_INCREMENT") + ddl.add(" " & col.name & " " & col.colType) + if col.isPk: ddl.add(" PRIMARY KEY") + if col.isNotNull: ddl.add(" NOT NULL") + if col.autoIncrement: ddl.add(" AUTO_INCREMENT") if col.fkTable.len > 0: - result.add(" REFERENCES " & col.fkTable & "(" & col.fkColumn & ")") - if i < tbl.columns.len - 1: result.add(",") - result.add("\n") + ddl.add(" REFERENCES " & col.fkTable & "(" & col.fkColumn & ")") + if i < tbl.columns.len - 1: ddl.add(",") + ddl.add("\n") # Sample data var kvPairs: seq[(string, seq[byte])] = @[] let rows = execScan(ctx, table) let sampleLimit = min(5, rows.len) if sampleLimit > 0: - result.add(");\n\n-- Sample data:\n") + ddl.add(");\n\n-- Sample data:\n") for i in 0.. 0: - result.add("\n-- Indexes: " & idxList.join(", ")) + ddl.add("\n-- Indexes: " & idxList.join(", ")) # RLS policies if table in ctx.policies and ctx.policies[table].len > 0: - result.add("\n-- RLS Policies:\n") + ddl.add("\n-- RLS Policies:\n") for pol in ctx.policies[table]: - result.add("-- CREATE POLICY " & pol.name & " FOR " & pol.command & "\n") + ddl.add("-- CREATE POLICY " & pol.name & " FOR " & pol.command & "\n") if tbl.foreignKeys.len > 0: - result.add("\n-- Foreign Keys:\n") + ddl.add("\n-- Foreign Keys:\n") for fk in tbl.foreignKeys: - result.add("-- " & fk.refTable & "(" & fk.refColumn & ") ON DELETE " & fk.onDelete & "\n") + ddl.add("-- " & fk.refTable & "(" & fk.refColumn & ") ON DELETE " & fk.onDelete & "\n") - return result + return ddl of "similarity_nodes": if expr.irFuncArgs.len < 1: return "[]" let graphName = evalExpr(expr.irFuncArgs[0], row, ctx) @@ -1508,7 +1553,6 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = finally: release(ctx.sharedLock.lock) return "0" - return "0" of "snowflake_id": # Snowflake ID: timestamp_ms(41 bits) | node_id(10 bits) | sequence(12 bits) var nodeId: int64 = 0 @@ -1662,20 +1706,30 @@ proc execInsert*(ctx: ExecutionContext, table: string, fields: seq[string], valu kvPairs: var seq[(string, seq[byte])]): int = if not hasPrivilege(ctx, table, "INSERT"): return 0 + let tblDef = if table in ctx.tables: ctx.tables[table] else: TableDef() var count = 0 for rowVals in values: var key = "" var keyFound = false var valParts: seq[string] = @[] + # Build composite PK key from all PK columns + if tblDef.pkColumns.len > 0: + var pkParts: seq[string] = @[] + for pkCol in tblDef.pkColumns: + let pkVal = getValue(rowVals, fields, pkCol) + pkParts.add(pkCol & "=" & escapeRowVal(pkVal)) + key = pkParts.join(":") + keyFound = true for i, f in fields: if i < rowVals.len: if not keyFound: key = f & "=" & escapeRowVal(rowVals[i]) keyFound = true - else: + elif tblDef.pkColumns.len == 0 or f.toLower() notin tblDef.pkColumns.mapIt(it.toLower()): valParts.add(f & "=" & escapeRowVal(rowVals[i])) elif f.len > 0: - valParts.add(f & "=") + if tblDef.pkColumns.len == 0 or f.toLower() notin tblDef.pkColumns.mapIt(it.toLower()): + valParts.add(f & "=") let valStr = valParts.join(",") let fullKey = table & "." & key @@ -2107,32 +2161,24 @@ proc validateConstraints*(ctx: ExecutionContext, tableName: string, if not typeOk: return (false, typeErr) - # FK check + # FK check — uses LSM get which searches memtable + SSTables if col.fkTable.len > 0 and col.fkColumn.len > 0 and not isNull(val): let fkKey = col.fkTable & "." & col.fkColumn & "=" & val let (fkExists, _) = ctx.db.get(fkKey) if not fkExists: - # Also check if value is in any row's first field - var found = false - let prefix = col.fkTable & "." - for entry in ctx.db.scanMemTable(): - if entry.deleted: continue - if entry.key.startsWith(prefix): - let rest = entry.key[prefix.len..^1] - if rest.startsWith(col.fkColumn & "=") and rest[col.fkColumn.len+1..^1] == val: - found = true - break - if not found: - return (false, "FOREIGN KEY violation: '" & val & "' not found in " & col.fkTable & "." & col.fkColumn) + return (false, "FOREIGN KEY violation: '" & val & "' not found in " & col.fkTable & "." & col.fkColumn) # PK uniqueness (skip during UPDATE — PK shouldn't change) if not skipPkCheck and tbl.pkColumns.len > 0: var pkVals: seq[string] = @[] + var pkParts: seq[string] = @[] for pkCol in tbl.pkColumns: - pkVals.add(getValue(rowVals, fields, pkCol)) + let pkVal = getValue(rowVals, fields, pkCol) + pkVals.add(pkVal) + pkParts.add(pkCol & "=" & escapeRowVal(pkVal)) let pkStr = pkVals.join("|") - # Check with field=value format (as stored by execInsert) - var pkKey = tableName & "." & tbl.pkColumns[0] & "=" & escapeRowVal(pkVals[0]) + # Check with composite PK format (as stored by execInsert) + let pkKey = tableName & "." & pkParts.join(":") let (exists, _) = ctx.db.get(pkKey) if exists: return (false, "UNIQUE constraint violated: duplicate key '" & pkStr & "'") @@ -4739,7 +4785,7 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu ctx.views[stmt.cvName] = stmt.cvQuery let viewKey = "_schema:views:" & stmt.cvName let viewSql = selectToSql(stmt.cvQuery) - let viewDdl = "CREATE VIEW " & stmt.cvName & " AS " & viewSql + let viewDdl = "CREATE VIEW \"" & sqlEscapeIdent(stmt.cvName) & "\" AS " & viewSql ctx.db.put(viewKey, cast[seq[byte]](viewDdl)) return okResult(msg="CREATE VIEW " & stmt.cvName) @@ -4762,7 +4808,7 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu ctx.tables[stmt.trigTable].triggers = triggers # Persist trigger to LSM-Tree let trigKey = "_schema:triggers:" & stmt.trigTable & ":" & stmt.trigName - let trigDdl = "CREATE TRIGGER " & stmt.trigName & " ON " & stmt.trigTable & " " & + let trigDdl = "CREATE TRIGGER \"" & sqlEscapeIdent(stmt.trigName) & "\" ON \"" & sqlEscapeIdent(stmt.trigTable) & "\" " & stmt.trigTiming & " " & stmt.trigEvent & " AS " & stmt.trigAction.strVal ctx.db.put(trigKey, cast[seq[byte]](trigDdl)) return okResult(msg="CREATE TRIGGER " & stmt.trigName) @@ -5031,7 +5077,7 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu ctx.users[stmt.cuName] = UserDef(name: stmt.cuName, passwordHash: stmt.cuPassword, isSuperuser: stmt.cuSuperuser, roles: @[]) let userKey = "_schema:users:" & stmt.cuName - let userDdl = "CREATE USER " & stmt.cuName & " WITH PASSWORD '" & stmt.cuPassword & "'" & + let userDdl = "CREATE USER \"" & sqlEscapeIdent(stmt.cuName) & "\" WITH PASSWORD '" & sqlEscapeString(stmt.cuPassword) & "'" & (if stmt.cuSuperuser: " SUPERUSER" else: " NOSUPERUSER") ctx.db.put(userKey, cast[seq[byte]](userDdl)) return okResult(msg="CREATE USER " & stmt.cuName) @@ -5050,7 +5096,7 @@ proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValu withCheckExpr: stmt.cpWithCheck)) ctx.policies[stmt.cpTable] = pols let polKey = "_schema:policies:" & stmt.cpTable & ":" & stmt.cpName - var polDdl = "CREATE POLICY " & stmt.cpName & " ON " & stmt.cpTable + var polDdl = "CREATE POLICY \"" & sqlEscapeIdent(stmt.cpName) & "\" ON \"" & sqlEscapeIdent(stmt.cpTable) & "\"" if stmt.cpCommand != "ALL": polDdl.add(" FOR " & stmt.cpCommand) if stmt.cpUsing != nil: diff --git a/src/barabadb/query/parser.nim b/src/barabadb/query/parser.nim index 74b1f62..2ce964e 100644 --- a/src/barabadb/query/parser.nim +++ b/src/barabadb/query/parser.nim @@ -181,8 +181,8 @@ proc parsePrimary(p: var Parser): Node = Node(kind: nkCase, caseExpr: caseExpr, caseWhens: whens, caseElse: elseExpr, line: tok.line, col: tok.col) else: - discard p.advance() - Node(kind: nkNullLit, line: tok.line, col: tok.col) + raise newException(ValueError, + "Unexpected token " & $tok.kind & " at line " & $tok.line & ", col " & $tok.col) proc parseOverClause(p: var Parser): Node = ## Parse OVER ( [PARTITION BY expr, ...] [ORDER BY expr [ASC|DESC], ...] [frame] ) @@ -324,7 +324,10 @@ proc parseComparison(p: var Parser): Node = if p.peek().kind == tkNot: discard p.advance() negated = true - discard p.advance() # consume NULL token (assumed) + if p.peek().kind != tkNull: + raise newException(ValueError, + "Expected NULL after IS " & (if negated: "NOT " else: "") & "at line " & $tok.line) + discard p.advance() # consume NULL token return Node(kind: nkIsExpr, isExpr: result, isNegated: negated, line: tok.line, col: tok.col) while p.peek().kind in {tkEq, tkNotEq, tkLt, tkLtEq, tkGt, tkGtEq, tkFtsMatch, tkDistanceOp, tkJsonContains, tkJsonContainedBy, tkJsonHasAny, tkJsonHasAll}: diff --git a/src/barabadb/storage/bloom.nim b/src/barabadb/storage/bloom.nim index 8b534de..0f2b038 100644 --- a/src/barabadb/storage/bloom.nim +++ b/src/barabadb/storage/bloom.nim @@ -81,7 +81,7 @@ proc deserialize*(bf: var BloomFilter, data: seq[byte]) = ) let numBytes = (bf.size + 7) div 8 if data.len < 8 + numBytes: - return + raise newException(ValueError, "Bloom filter data too short: expected " & $(8 + numBytes) & " bytes, got " & $data.len) for i in 0.. 0: + db.flushUnsafe() + db.immutableMem = db.memTable + db.memTable = newMemTable(db.memMaxSize) + discard db.memTable.put(key, value, ts, deleted) + +proc deleteUnsafe*(db: LSMTree, key: string) = + putUnsafe(db, key, @[], deleted = true) + proc getUnsafe(db: LSMTree, key: string): (bool, seq[byte]) = let (found, entry) = db.memTable.get(key) if found: @@ -478,6 +493,9 @@ proc flush*(db: LSMTree) = proc close*(db: LSMTree) = acquire(db.lock) defer: release(db.lock) + # Flush both memtables to avoid data loss + while db.immutableMem.len > 0: + flushUnsafe(db) flushUnsafe(db) for sst in db.sstables.mitems: sst.close() @@ -531,7 +549,8 @@ proc scanAll*(db: LSMTree): seq[(string, seq[byte])] = result.add((e.key, e.value)) # Scan SSTables from newest to oldest - for sst in db.sstables: + for i in countdown(db.sstables.high, db.sstables.low): + let sst = db.sstables[i] for key, offset in sst.index: if key notin seen: seen[key] = true diff --git a/src/barabadb/storage/mmap.nim b/src/barabadb/storage/mmap.nim index 477b7e8..bc16656 100644 --- a/src/barabadb/storage/mmap.nim +++ b/src/barabadb/storage/mmap.nim @@ -54,14 +54,18 @@ proc openMmap*(path: string, mode: MmapMode = mmReadOnly): MmapFile = let mapped = mmap(nil, fileSize, prot, flags, fd, 0) if mapped == MAP_FAILED: - discard close(fd) + discard posix.close(fd) return MmapFile(path: path, regions: @[], totalSize: 0, pageSize: PageSize) + # Kernel keeps the mapping alive independently of the fd; close it now + # to avoid leaking one fd per SSTable. + discard posix.close(fd) + let region = MmapRegion( data: cast[ptr UncheckedArray[byte]](mapped), size: fileSize, offset: 0, - fd: fd, + fd: -1, mode: mode, ) @@ -116,16 +120,21 @@ proc adviseRandom*(mf: MmapFile) = proc adviseWillNeed*(mf: MmapFile, offset: int, size: int) = if mf.regions.len > 0: + if offset < 0 or offset + size > mf.regions[0].size: + return discard madvise(addr mf.regions[0].data[offset], size, MADV_WILLNEED) proc adviseDontNeed*(mf: MmapFile, offset: int, size: int) = if mf.regions.len > 0: + if offset < 0 or offset + size > mf.regions[0].size: + return discard madvise(addr mf.regions[0].data[offset], size, MADV_DONTNEED) proc close*(mf: MmapFile) = for region in mf.regions: discard munmap(region.data, region.size) - discard close(cint(region.fd)) + if region.fd != -1: + discard close(cint(region.fd)) mf.regions.setLen(0) proc size*(mf: MmapFile): int = mf.totalSize diff --git a/src/barabadb/storage/recovery.nim b/src/barabadb/storage/recovery.nim index e74950b..daed440 100644 --- a/src/barabadb/storage/recovery.nim +++ b/src/barabadb/storage/recovery.nim @@ -32,6 +32,7 @@ type dataDir*: string entries*: seq[RecoveredEntry] result*: RecoveryResult + lastTxnId*: uint64 # tracks commits seen in WAL proc newCrashRecovery*(walDir: string, dataDir: string): CrashRecovery = CrashRecovery( @@ -101,6 +102,7 @@ proc scanWAL*(rec: CrashRecovery): seq[RecoveredEntry] = discard stream.close() + rec.lastTxnId = txnId proc analyze*(rec: CrashRecovery): RecoveryResult = rec.entries = rec.scanWAL() @@ -109,10 +111,7 @@ proc analyze*(rec: CrashRecovery): RecoveryResult = rec.result = RecoveryResult(state: recDone, applied: false) return rec.result - var lastCommitted: uint64 = 0 - for entry in rec.entries: - if entry.txnId > lastCommitted: - lastCommitted = entry.txnId + var lastCommitted = rec.lastTxnId var redoCount = 0 var undoCount = 0 @@ -149,11 +148,11 @@ proc recover*(rec: CrashRecovery, db: LSMTree = nil): RecoveryResult = var undoCount = 0 for entry in rec.entries: if entry.txnId < analysis.lastTxn: - # Committed — redo + # Committed — redo (bypass WAL to avoid duplicate entries) if entry.isDelete: - db.delete(entry.key) + db.deleteUnsafe(entry.key) else: - db.put(entry.key, entry.value) + db.putUnsafe(entry.key, entry.value) inc redoCount else: # Uncommitted — skip (undo) diff --git a/src/barabadb/storage/wal.nim b/src/barabadb/storage/wal.nim index 87ea0d5..dbe26bb 100644 --- a/src/barabadb/storage/wal.nim +++ b/src/barabadb/storage/wal.nim @@ -30,10 +30,11 @@ proc newWriteAheadLog*(dir: string, syncOnWrite: bool = true): WriteAheadLog = createDir(dir) let path = dir / "wal.log" let exists = fileExists(path) + let isEmpty = if exists: getFileSize(path) == 0 else: true let stream = if exists: newFileStream(path, fmAppend) else: newFileStream(path, fmWrite) if stream == nil: raise newException(IOError, "Cannot open WAL: " & path) - if not exists: + if not exists or isEmpty: stream.write(WALMagic) stream.write(WALVersion) stream.flush() @@ -78,14 +79,17 @@ proc writeCommit*(wal: var WriteAheadLog, timestamp: uint64) = proc sync*(wal: var WriteAheadLog) = wal.stream.flush() - let fd = posix.open(cstring(wal.path), O_RDONLY) + # Re-open with O_RDWR so fsync operates on a write-capable fd. + # Not ideal (two fds for same file) but avoids accessing private + # FileStream internals that vary across Nim versions. + let fd = posix.open(cstring(wal.path), O_RDWR) if fd != -1: discard posix.fsync(fd) discard posix.close(fd) proc close*(wal: var WriteAheadLog) = wal.stream.flush() - let fd = posix.open(cstring(wal.path), O_RDONLY) + let fd = posix.open(cstring(wal.path), O_RDWR) if fd != -1: discard posix.fsync(fd) discard posix.close(fd) @@ -100,10 +104,10 @@ proc readEntries*(walPath: string, untilTimestamp: uint64 = 0): seq[WalEntry] = let s = newFileStream(walPath, fmRead) if s == nil: return # Skip header - var magic: uint32 - var version: uint32 - discard s.readData(addr magic, 4) - discard s.readData(addr version, 4) + var magic: uint32 = 0 + var version: uint32 = 0 + if s.readData(addr magic, 4) != 4: return + if s.readData(addr version, 4) != 4: return if magic != WALMagic: return while not s.atEnd: var kind: uint8 diff --git a/tests/fuzz_test b/tests/fuzz_test new file mode 100755 index 0000000..4214575 Binary files /dev/null and b/tests/fuzz_test differ diff --git a/tests/nimforum_smoke_test b/tests/nimforum_smoke_test new file mode 100755 index 0000000..64c172f Binary files /dev/null and b/tests/nimforum_smoke_test differ diff --git a/tests/prop_test b/tests/prop_test new file mode 100755 index 0000000..c5cd712 Binary files /dev/null and b/tests/prop_test differ diff --git a/tests/test_all.nim b/tests/test_all.nim index d4724e3..c9d8492 100644 --- a/tests/test_all.nim +++ b/tests/test_all.nim @@ -3845,6 +3845,59 @@ suite "Foreign Key Enforcement": check childSel.rows.len == 0 +suite "Composite Primary Key": + var db: LSMTree + var ctx: qexec.ExecutionContext + var tmpDir: string + + setup: + tmpDir = getTempDir() / "baradb_compk_test_" & $getMonoTime().ticks + db = newLSMTree(tmpDir) + ctx = qexec.newExecutionContext(db) + + teardown: + removeDir(tmpDir) + + test "Composite PK insert stores distinct rows": + discard qexec.executeQuery(ctx, parse("CREATE TABLE orders (user_id INT, order_id INT, amount INT, PRIMARY KEY (user_id, order_id))")) + discard qexec.executeQuery(ctx, parse("INSERT INTO orders (user_id, order_id, amount) VALUES (1, 100, 50)")) + discard qexec.executeQuery(ctx, parse("INSERT INTO orders (user_id, order_id, amount) VALUES (1, 101, 75)")) + let sel = qexec.executeQuery(ctx, parse("SELECT * FROM orders")) + check sel.success + check sel.rows.len == 2 + + test "Composite PK prevents duplicate key": + discard qexec.executeQuery(ctx, parse("CREATE TABLE items (warehouse INT, sku TEXT, qty INT, PRIMARY KEY (warehouse, sku))")) + discard qexec.executeQuery(ctx, parse("INSERT INTO items (warehouse, sku, qty) VALUES (1, 'ABC', 10)")) + let dup = qexec.executeQuery(ctx, parse("INSERT INTO items (warehouse, sku, qty) VALUES (1, 'ABC', 20)")) + check not dup.success + check dup.message.contains("UNIQUE constraint violated") + + test "Composite PK allows same first column with different second": + discard qexec.executeQuery(ctx, parse("CREATE TABLE scores (player INT, level INT, score INT, PRIMARY KEY (player, level))")) + discard qexec.executeQuery(ctx, parse("INSERT INTO scores (player, level, score) VALUES (1, 1, 100)")) + discard qexec.executeQuery(ctx, parse("INSERT INTO scores (player, level, score) VALUES (1, 2, 200)")) + let sel = qexec.executeQuery(ctx, parse("SELECT * FROM scores")) + check sel.success + check sel.rows.len == 2 + + test "Composite PK SELECT returns correct rows": + discard qexec.executeQuery(ctx, parse("CREATE TABLE pair (a INT, b INT, val TEXT, PRIMARY KEY (a, b))")) + discard qexec.executeQuery(ctx, parse("INSERT INTO pair (a, b, val) VALUES (1, 2, 'first')")) + discard qexec.executeQuery(ctx, parse("INSERT INTO pair (a, b, val) VALUES (1, 3, 'second')")) + let sel = qexec.executeQuery(ctx, parse("SELECT * FROM pair")) + check sel.success + check sel.rows.len == 2 + # Verify both rows exist (order may vary) + var foundFirst = false + var foundSecond = false + for row in sel.rows: + if row["val"] == "first": foundFirst = true + if row["val"] == "second": foundSecond = true + check foundFirst + check foundSecond + + suite "Hybrid RAG Search": var db: LSMTree var ctx: qexec.ExecutionContext