diff --git a/clients/javascript/baradb.js b/clients/javascript/baradb.js index 265165b..c963a53 100644 --- a/clients/javascript/baradb.js +++ b/clients/javascript/baradb.js @@ -3,6 +3,7 @@ * * Binary protocol client for BaraDB database. * Communicates via the BaraDB Wire Protocol (binary, big-endian). + * Requires Node.js (uses 'net' module for TCP). * * Install: * npm install baradb @@ -16,23 +17,42 @@ * console.log(row.name); * } * await client.close(); + * + * Parameterized Queries: + * const result = await client.queryParams( + * 'SELECT * FROM users WHERE age > $1', + * [WireValue.int64(18)] + * ); + * + * Authentication: + * await client.auth('jwt-token-here'); */ -const MsgKind = { +const net = require('net'); + +const MsgKind = Object.freeze({ + CLIENT_HANDSHAKE: 0x01, QUERY: 0x02, + QUERY_PARAMS: 0x03, + EXECUTE: 0x04, BATCH: 0x05, TRANSACTION: 0x06, CLOSE: 0x07, PING: 0x08, AUTH: 0x09, + SERVER_HANDSHAKE: 0x80, READY: 0x81, DATA: 0x82, COMPLETE: 0x83, ERROR: 0x84, + AUTH_CHALLENGE: 0x85, AUTH_OK: 0x86, -}; + SCHEMA_CHANGE: 0x87, + PONG: 0x88, + TRANSACTION_STATE: 0x89, +}); -const FieldKind = { +const FieldKind = Object.freeze({ NULL: 0x00, BOOL: 0x01, INT8: 0x02, @@ -47,13 +67,136 @@ const FieldKind = { OBJECT: 0x0B, VECTOR: 0x0C, JSON: 0x0D, -}; +}); -const ResultFormat = { +const ResultFormat = Object.freeze({ BINARY: 0x00, JSON: 0x01, TEXT: 0x02, -}; +}); + +class WireValue { + constructor(kind, value) { + this.kind = kind; + this.value = value; + } + + static null() { return new WireValue(FieldKind.NULL); } + static bool(val) { return new WireValue(FieldKind.BOOL, val); } + static int8(val) { return new WireValue(FieldKind.INT8, val); } + static int16(val) { return new WireValue(FieldKind.INT16, val); } + static int32(val) { return new WireValue(FieldKind.INT32, val); } + static int64(val) { return new WireValue(FieldKind.INT64, val); } + static float32(val) { return new WireValue(FieldKind.FLOAT32, val); } + static float64(val) { return new WireValue(FieldKind.FLOAT64, val); } + static string(val) { return new WireValue(FieldKind.STRING, val); } + static bytes(val) { return new WireValue(FieldKind.BYTES, val); } + static array(val) { return new WireValue(FieldKind.ARRAY, val); } + static object(val) { return new WireValue(FieldKind.OBJECT, val); } + static vector(val) { return new WireValue(FieldKind.VECTOR, val); } + static json(val) { return new WireValue(FieldKind.JSON, val); } + + serialize() { + const parts = [Buffer.from([this.kind])]; + switch (this.kind) { + case FieldKind.NULL: + break; + case FieldKind.BOOL: + parts.push(Buffer.from([this.value ? 1 : 0])); + break; + case FieldKind.INT8: { + const b = Buffer.alloc(1); + b.writeInt8(this.value, 0); + parts.push(b); + break; + } + case FieldKind.INT16: { + const b = Buffer.alloc(2); + b.writeInt16BE(this.value, 0); + parts.push(b); + break; + } + case FieldKind.INT32: { + const b = Buffer.alloc(4); + b.writeInt32BE(this.value, 0); + parts.push(b); + break; + } + case FieldKind.INT64: { + const b = Buffer.alloc(8); + b.writeBigInt64BE(BigInt(this.value), 0); + parts.push(b); + break; + } + case FieldKind.FLOAT32: { + const b = Buffer.alloc(4); + b.writeFloatBE(this.value, 0); + parts.push(b); + break; + } + case FieldKind.FLOAT64: { + const b = Buffer.alloc(8); + b.writeDoubleBE(this.value, 0); + parts.push(b); + break; + } + case FieldKind.STRING: { + const strBuf = Buffer.from(this.value, 'utf-8'); + const len = Buffer.alloc(4); + len.writeUInt32BE(strBuf.length, 0); + parts.push(len, strBuf); + break; + } + case FieldKind.BYTES: { + const len = Buffer.alloc(4); + len.writeUInt32BE(this.value.length, 0); + parts.push(len, this.value); + break; + } + case FieldKind.ARRAY: { + const count = Buffer.alloc(4); + count.writeUInt32BE(this.value.length, 0); + parts.push(count); + for (const item of this.value) { + parts.push(item.serialize()); + } + break; + } + case FieldKind.OBJECT: { + const entries = Object.entries(this.value); + const count = Buffer.alloc(4); + count.writeUInt32BE(entries.length, 0); + parts.push(count); + for (const [key, val] of entries) { + const keyBuf = Buffer.from(key, 'utf-8'); + const keyLen = Buffer.alloc(4); + keyLen.writeUInt32BE(keyBuf.length, 0); + parts.push(keyLen, keyBuf, val.serialize()); + } + break; + } + case FieldKind.VECTOR: { + const dim = Buffer.alloc(4); + dim.writeUInt32BE(this.value.length, 0); + parts.push(dim); + for (const f of this.value) { + const fb = Buffer.alloc(4); + fb.writeFloatBE(f, 0); + parts.push(fb); + } + break; + } + case FieldKind.JSON: { + const strBuf = Buffer.from(this.value, 'utf-8'); + const len = Buffer.alloc(4); + len.writeUInt32BE(strBuf.length, 0); + parts.push(len, strBuf); + break; + } + } + return Buffer.concat(parts); + } +} class QueryResult { constructor() { @@ -99,17 +242,50 @@ class Client { this.socket = null; this.connected = false; this.requestId = 0; + this._buffer = Buffer.alloc(0); + this._pendingResolve = null; } - /** Connect to the BaraDB server */ async connect() { - // In Node.js: const net = require('net'); - // In browser: WebSocket connection - this.socket = null; // net.connect(this.port, this.host); - this.connected = true; + return new Promise((resolve, reject) => { + this.socket = net.createConnection({ host: this.host, port: this.port }); + this.socket.setTimeout(this.timeout); + + this.socket.on('connect', () => { + this.connected = true; + resolve(); + }); + + this.socket.on('data', (data) => { + this._buffer = Buffer.concat([this._buffer, data]); + if (this._pendingResolve) { + this._pendingResolve(); + } + }); + + this.socket.on('error', (err) => { + this.connected = false; + reject(err); + }); + + this.socket.on('close', () => { + this.connected = false; + }); + + this.socket.on('timeout', () => { + this.socket.destroy(); + this.connected = false; + }); + }); } async close() { + if (this.socket && this.connected) { + try { + const msg = this._build(MsgKind.CLOSE, Buffer.alloc(0)); + this.socket.write(msg); + } catch (_) {} + } if (this.socket) { this.socket.destroy(); this.socket = null; @@ -125,75 +301,88 @@ class Client { return ++this.requestId; } - /** Execute a BaraQL query */ - async query(sql) { - const encoder = new TextEncoder(); - const queryBytes = encoder.encode(sql); - const payload = new Uint8Array(4 + queryBytes.length + 1); - const view = new DataView(payload.buffer); - view.setUint32(0, queryBytes.length, false); // big-endian - payload.set(queryBytes, 4); - payload[4 + queryBytes.length] = ResultFormat.BINARY; - - const msg = this._build(MsgKind.QUERY, payload); - // await this.socket.write(msg); - // const header = await this.socket.read(12); - // ... parse response - - return new QueryResult(); + async _recvExact(size) { + while (this._buffer.length < size) { + await new Promise((resolve, reject) => { + this._pendingResolve = resolve; + const timer = setTimeout(() => { + this._pendingResolve = null; + reject(new Error('Receive timeout')); + }, this.timeout); + this._pendingResolve = () => { + clearTimeout(timer); + this._pendingResolve = null; + resolve(); + }; + }); + } + const data = this._buffer.subarray(0, size); + this._buffer = this._buffer.subarray(size); + return data; } - /** Parse server response from header + payload buffers */ - _parseResponse(header, payload) { - const view = new DataView(header.buffer); - const kind = view.getUint32(0, false); - const len = view.getUint32(4, false); - const reqId = view.getUint32(8, false); + async _readHeader() { + const header = await this._recvExact(12); + return { + kind: header.readUInt32BE(0), + length: header.readUInt32BE(4), + requestId: header.readUInt32BE(8), + }; + } + async _readError(length) { + const data = await this._recvExact(length); + const code = data.readUInt32BE(0); + const msgLen = data.readUInt32BE(4); + const msg = data.subarray(8, 8 + msgLen).toString('utf-8'); + return new Error(`BaraDB error ${code}: ${msg}`); + } + + async _readDataResponse(length) { + const data = await this._recvExact(length); const result = new QueryResult(); + let pos = 0; - if (kind === MsgKind.ERROR && payload.length >= 8) { - const code = new DataView(payload.buffer).getUint32(0, false); - const msgLen = new DataView(payload.buffer).getUint32(4, false); - const msg = new TextDecoder().decode(payload.slice(8, 8 + msgLen)); - throw new Error(`BaraDB error ${code}: ${msg}`); + const colCount = data.readUInt32BE(pos); + pos += 4; + + const cols = []; + for (let i = 0; i < colCount; i++) { + const sLen = data.readUInt32BE(pos); + pos += 4; + cols.push(data.subarray(pos, pos + sLen).toString('utf-8')); + pos += sLen; } - if (kind === MsgKind.DATA) { - let pos = 0; - const colCount = new DataView(payload.buffer).getUint32(pos, false); - pos += 4; - const cols = []; - for (let i = 0; i < colCount; i++) { - const sLen = new DataView(payload.buffer).getUint32(pos, false); - pos += 4; - cols.push(new TextDecoder().decode(payload.slice(pos, pos + sLen))); - pos += sLen; - } - result.columns = cols; - const colTypes = []; - for (let i = 0; i < colCount; i++) { - colTypes.push(payload[pos]); - pos++; - } - result.columnTypes = colTypes; - const rowCount = new DataView(payload.buffer).getUint32(pos, false); - pos += 4; - for (let r = 0; r < rowCount; r++) { - const row = []; - for (let c = 0; c < colCount; c++) { - row.push(this._readValue(payload, pos)); - pos = this._lastReadPos; - } - result.rows.push(row); - } - result.rowCount = rowCount; - // COMPLETE message should follow - caller reads it separately - return result; + const colTypes = []; + for (let i = 0; i < colCount; i++) { + colTypes.push(data[pos]); + pos++; } - if (kind === MsgKind.COMPLETE && payload.length >= 4) { - result.affectedRows = new DataView(payload.buffer).getUint32(0, false); + const rowCount = data.readUInt32BE(pos); + pos += 4; + + for (let r = 0; r < rowCount; r++) { + const row = []; + for (let c = 0; c < colCount; c++) { + const { value, newPos } = this._readValue(data, pos); + row.push(value); + pos = newPos; + } + result.rows.push(row); + } + + result.columns = cols; + result.columnTypes = colTypes; + result.rowCount = rowCount; + + const comp = await this._readHeader(); + if (comp.kind === MsgKind.COMPLETE) { + const compData = await this._recvExact(comp.length); + result.affectedRows = compData.readUInt32BE(0); + } else if (comp.kind === MsgKind.ERROR) { + throw await this._readError(comp.length); } return result; @@ -206,74 +395,155 @@ class Client { switch (kind) { case FieldKind.NULL: result = null; break; case FieldKind.BOOL: result = payload[pos] !== 0; pos++; break; - case FieldKind.INT8: result = payload[pos] << 24 >> 24; pos++; break; - case FieldKind.INT16: result = new DataView(payload.buffer).getInt16(pos, false); pos += 2; break; - case FieldKind.INT32: result = new DataView(payload.buffer).getInt32(pos, false); pos += 4; break; - case FieldKind.INT64: result = new DataView(payload.buffer).getBigInt64(pos, false); pos += 8; break; - case FieldKind.FLOAT32: result = new DataView(payload.buffer).getFloat32(pos, false); pos += 4; break; - case FieldKind.FLOAT64: result = new DataView(payload.buffer).getFloat64(pos, false); pos += 8; break; + case FieldKind.INT8: result = payload.readInt8(pos); pos++; break; + case FieldKind.INT16: result = payload.readInt16BE(pos); pos += 2; break; + case FieldKind.INT32: result = payload.readInt32BE(pos); pos += 4; break; + case FieldKind.INT64: result = payload.readBigInt64BE(pos); pos += 8; break; + case FieldKind.FLOAT32: result = payload.readFloatBE(pos); pos += 4; break; + case FieldKind.FLOAT64: result = payload.readDoubleBE(pos); pos += 8; break; case FieldKind.STRING: { - const len = new DataView(payload.buffer).getUint32(pos, false); + const len = payload.readUInt32BE(pos); pos += 4; - result = new TextDecoder().decode(payload.slice(pos, pos + len)); + result = payload.subarray(pos, pos + len).toString('utf-8'); pos += len; break; } case FieldKind.BYTES: { - const len = new DataView(payload.buffer).getUint32(pos, false); + const len = payload.readUInt32BE(pos); pos += 4; - result = payload.slice(pos, pos + len); + result = payload.subarray(pos, pos + len); pos += len; break; } case FieldKind.ARRAY: { - const count = new DataView(payload.buffer).getUint32(pos, false); + const count = payload.readUInt32BE(pos); pos += 4; result = []; for (let i = 0; i < count; i++) { - result.push(this._readValue(payload, pos)); - pos = this._lastReadPos; + const { value, newPos } = this._readValue(payload, pos); + result.push(value); + pos = newPos; } break; } case FieldKind.OBJECT: { - const count = new DataView(payload.buffer).getUint32(pos, false); + const count = payload.readUInt32BE(pos); pos += 4; result = {}; for (let i = 0; i < count; i++) { - const kLen = new DataView(payload.buffer).getUint32(pos, false); + const kLen = payload.readUInt32BE(pos); pos += 4; - const key = new TextDecoder().decode(payload.slice(pos, pos + kLen)); + const key = payload.subarray(pos, pos + kLen).toString('utf-8'); pos += kLen; - result[key] = this._readValue(payload, pos); - pos = this._lastReadPos; + const { value, newPos } = this._readValue(payload, pos); + result[key] = value; + pos = newPos; } break; } case FieldKind.VECTOR: { - const dim = new DataView(payload.buffer).getUint32(pos, false); + const dim = payload.readUInt32BE(pos); pos += 4; result = []; for (let i = 0; i < dim; i++) { - result.push(new DataView(payload.buffer).getFloat32(pos, false)); + result.push(payload.readFloatBE(pos)); pos += 4; } break; } case FieldKind.JSON: { - const len = new DataView(payload.buffer).getUint32(pos, false); + const len = payload.readUInt32BE(pos); pos += 4; - result = new TextDecoder().decode(payload.slice(pos, pos + len)); + result = payload.subarray(pos, pos + len).toString('utf-8'); pos += len; break; } default: result = null; } - this._lastReadPos = pos; - return result; + return { value: result, newPos: pos }; + } + + async auth(token) { + if (!this.connected) throw new Error('Not connected'); + const tokenBuf = Buffer.from(token, 'utf-8'); + const payload = Buffer.alloc(4 + tokenBuf.length); + payload.writeUInt32BE(tokenBuf.length, 0); + tokenBuf.copy(payload, 4); + const msg = this._build(MsgKind.AUTH, payload); + this.socket.write(msg); + + const header = await this._readHeader(); + if (header.kind === MsgKind.AUTH_OK) return; + if (header.kind === MsgKind.ERROR) throw await this._readError(header.length); + throw new Error(`Unexpected auth response: 0x${header.kind.toString(16)}`); + } + + async ping() { + if (!this.connected) throw new Error('Not connected'); + const msg = this._build(MsgKind.PING, Buffer.alloc(0)); + this.socket.write(msg); + + const header = await this._readHeader(); + if (header.kind === MsgKind.PONG) return true; + if (header.kind === MsgKind.ERROR) throw await this._readError(header.length); + return false; + } + + async query(sql) { + if (!this.connected) throw new Error('Not connected'); + const queryBuf = Buffer.from(sql, 'utf-8'); + const payload = Buffer.alloc(4 + queryBuf.length + 1); + payload.writeUInt32BE(queryBuf.length, 0); + queryBuf.copy(payload, 4); + payload[4 + queryBuf.length] = ResultFormat.BINARY; + + const msg = this._build(MsgKind.QUERY, payload); + this.socket.write(msg); + + const header = await this._readHeader(); + if (header.kind === MsgKind.ERROR) throw await this._readError(header.length); + if (header.kind === MsgKind.DATA) return await this._readDataResponse(header.length); + if (header.kind === MsgKind.COMPLETE) { + const data = await this._recvExact(header.length); + const result = new QueryResult(); + result.affectedRows = data.readUInt32BE(0); + return result; + } + return new QueryResult(); + } + + async queryParams(sql, params = []) { + if (!this.connected) throw new Error('Not connected'); + const queryBuf = Buffer.from(sql, 'utf-8'); + const paramParts = []; + for (const p of params) { + paramParts.push(p.serialize()); + } + const paramsBuf = Buffer.concat(paramParts); + + const payload = Buffer.alloc(4 + queryBuf.length + 1 + 4 + paramsBuf.length); + let pos = 0; + payload.writeUInt32BE(queryBuf.length, pos); pos += 4; + queryBuf.copy(payload, pos); pos += queryBuf.length; + payload[pos] = ResultFormat.BINARY; pos++; + payload.writeUInt32BE(params.length, pos); pos += 4; + paramsBuf.copy(payload, pos); + + const msg = this._build(MsgKind.QUERY_PARAMS, payload); + this.socket.write(msg); + + const header = await this._readHeader(); + if (header.kind === MsgKind.ERROR) throw await this._readError(header.length); + if (header.kind === MsgKind.DATA) return await this._readDataResponse(header.length); + if (header.kind === MsgKind.COMPLETE) { + const data = await this._recvExact(header.length); + const result = new QueryResult(); + result.affectedRows = data.readUInt32BE(0); + return result; + } + return new QueryResult(); } - /** Execute a BaraQL statement (INSERT, UPDATE, DELETE) */ async execute(sql) { const result = await this.query(sql); return result.affectedRows; @@ -281,16 +551,11 @@ class Client { _build(kind, payload) { const reqId = this._nextId(); - const header = new ArrayBuffer(12); - const view = new DataView(header); - view.setUint32(0, kind, false); - view.setUint32(4, payload.length, false); - view.setUint32(8, reqId, false); - - const msg = new Uint8Array(12 + payload.length); - msg.set(new Uint8Array(header), 0); - msg.set(payload, 12); - return msg; + const header = Buffer.alloc(12); + header.writeUInt32BE(kind, 0); + header.writeUInt32BE(payload.length, 4); + header.writeUInt32BE(reqId, 8); + return Buffer.concat([header, payload]); } } @@ -376,4 +641,4 @@ class QueryBuilder { } } -module.exports = { Client, QueryBuilder, QueryResult, MsgKind, FieldKind, ResultFormat }; +module.exports = { Client, QueryBuilder, QueryResult, WireValue, MsgKind, FieldKind, ResultFormat }; diff --git a/clients/nim/src/baradb/client.nim b/clients/nim/src/baradb/client.nim index 723f38a..96834d2 100644 --- a/clients/nim/src/baradb/client.nim +++ b/clients/nim/src/baradb/client.nim @@ -30,17 +30,27 @@ type fkJson = 0x0D MsgKind* = enum + # Client messages mkClientHandshake = 0x01 mkQuery = 0x02 + mkQueryParams = 0x03 + mkExecute = 0x04 mkBatch = 0x05 + mkTransaction = 0x06 + mkClose = 0x07 mkPing = 0x08 mkAuth = 0x09 - # Server + # Server messages + mkServerHandshake = 0x80 mkReady = 0x81 mkData = 0x82 mkComplete = 0x83 mkError = 0x84 + mkAuthChallenge = 0x85 + mkAuthOk = 0x86 + mkSchemaChange = 0x87 mkPong = 0x88 + mkTransactionState = 0x89 ResultFormat* = enum rfBinary = 0x00 @@ -69,6 +79,11 @@ proc writeUint32(buf: var seq[byte], val: uint32) = bigEndian32(addr bytes, unsafeAddr val) buf.add(bytes) +proc writeUint64(buf: var seq[byte], val: uint64) = + var bytes: array[8, byte] + bigEndian64(addr bytes, unsafeAddr val) + buf.add(bytes) + proc writeString(buf: var seq[byte], s: string) = buf.writeUint32(uint32(s.len)) for ch in s: @@ -80,6 +95,12 @@ proc readUint32(buf: openArray[byte], pos: var int): uint32 = bigEndian32(addr result, unsafeAddr bytes) pos += 4 +proc readUint64(buf: openArray[byte], pos: var int): uint64 = + var bytes: array[8, byte] + for i in 0..7: bytes[i] = buf[pos + i] + bigEndian64(addr result, unsafeAddr bytes) + pos += 8 + proc readString(buf: openArray[byte], pos: var int): string = let len = int(readUint32(buf, pos)) result = newString(len) @@ -98,10 +119,7 @@ proc serializeValue*(buf: var seq[byte], val: WireValue) = bigEndian16(addr bytes16, unsafeAddr val.int16Val) buf.add(bytes16) of fkInt32: buf.writeUint32(uint32(val.int32Val)) - of fkInt64: - var bytes: array[8, byte] - bigEndian64(addr bytes, unsafeAddr val.int64Val) - buf.add(bytes) + of fkInt64: buf.writeUint64(uint64(val.int64Val)) of fkFloat32: var bytes32: array[4, byte] copyMem(addr bytes32, unsafeAddr val.float32Val, 4) @@ -129,7 +147,7 @@ proc serializeValue*(buf: var seq[byte], val: WireValue) = var fb: array[4, byte] copyMem(addr fb, unsafeAddr f, 4) buf.add(fb) - else: discard + of fkJson: buf.writeString(val.jsonVal) proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue = let kind = FieldKind(buf[pos]) @@ -152,12 +170,7 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue = of fkInt32: result = WireValue(kind: fkInt32, int32Val: int32(readUint32(buf, pos))) of fkInt64: - var bytes: array[8, byte] - for i in 0..7: bytes[i] = buf[pos + i] - var v: int64 - bigEndian64(addr v, unsafeAddr bytes) - result = WireValue(kind: fkInt64, int64Val: v) - pos += 8 + result = WireValue(kind: fkInt64, int64Val: int64(readUint64(buf, pos))) of fkFloat32: var v32: float32 copyMem(addr v32, addr buf[pos], 4) @@ -202,8 +215,6 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue = result = WireValue(kind: fkVector, vecVal: vec) of fkJson: result = WireValue(kind: fkJson, jsonVal: readString(buf, pos)) - else: - result = WireValue(kind: fkNull) proc buildMessage*(kind: MsgKind, requestId: uint32, payload: seq[byte]): seq[byte] = result = @[] @@ -218,6 +229,20 @@ proc makeQueryMessage*(requestId: uint32, query: string): seq[byte] = payload.add(byte(rfBinary)) buildMessage(mkQuery, requestId, payload) +proc makeQueryParamsMessage*(requestId: uint32, query: string, params: seq[WireValue]): seq[byte] = + var payload: seq[byte] = @[] + payload.writeString(query) + payload.add(byte(rfBinary)) + payload.writeUint32(uint32(params.len)) + for p in params: + payload.serializeValue(p) + buildMessage(mkQueryParams, requestId, payload) + +proc makeAuthMessage*(requestId: uint32, token: string): seq[byte] = + var payload: seq[byte] = @[] + payload.writeString(token) + buildMessage(mkAuth, requestId, payload) + # === Client Library === type @@ -257,16 +282,20 @@ proc connect*(client: BaraClient) {.async.} = await client.socket.connect(client.config.host, Port(client.config.port)) client.connected = true +proc nextId(client: BaraClient): uint32 = + inc client.requestId; client.requestId + proc close*(client: BaraClient) = if client.connected: + try: + let msg = buildMessage(mkClose, client.nextId(), @[]) + waitFor client.socket.send(cast[string](msg)) + except: discard client.socket.close() client.connected = false proc isConnected*(client: BaraClient): bool = client.connected -proc nextId(client: BaraClient): uint32 = - inc client.requestId; client.requestId - proc wireValueToString*(wv: WireValue): string = case wv.kind of fkNull: return "" @@ -367,11 +396,48 @@ proc exec*(client: BaraClient, sql: string): Future[int] {.async.} = let qr = await client.query(sql) return qr.affectedRows -proc ping*(client: BaraClient) {.async.} = - if not client.connected: return +proc auth*(client: BaraClient, token: string) {.async.} = + if not client.connected: + raise newException(IOError, "Not connected") + + let msg = makeAuthMessage(client.nextId(), token) + await client.socket.send(cast[string](msg)) + + let headerData = await client.socket.recv(12) + if headerData.len < 12: + raise newException(IOError, "Connection closed") + + var pos = 0 + let hdrData = cast[seq[byte]](headerData) + let kind = MsgKind(readUint32(hdrData, pos)) + let payloadLen = int(readUint32(hdrData, pos)) + discard readUint32(hdrData, pos) + + if kind == mkAuthOk: + return + elif kind == mkError: + let payloadStr = await client.socket.recv(payloadLen) + var epos = 0 + let emsg = readString(cast[seq[byte]](payloadStr), epos) + raise newException(IOError, "Auth failed: " & emsg) + else: + raise newException(IOError, "Unexpected auth response: 0x" & toHex(uint32(kind), 2)) + +proc ping*(client: BaraClient): Future[bool] {.async.} = + if not client.connected: + return false let msg = buildMessage(mkPing, client.nextId(), @[]) await client.socket.send(cast[string](msg)) + let headerData = await client.socket.recv(12) + if headerData.len < 12: + return false + + var pos = 0 + let hdrData = cast[seq[byte]](headerData) + let kind = MsgKind(readUint32(hdrData, pos)) + return kind == mkPong + # === Fluent Query Builder === type @@ -468,6 +534,15 @@ proc close*(client: SyncClient) = proc query*(client: SyncClient, sql: string): QueryResult = waitFor client.asyncClient.query(sql) +proc query*(client: SyncClient, sql: string, params: seq[WireValue]): QueryResult = + waitFor client.asyncClient.query(sql, params) + +proc auth*(client: SyncClient, token: string) = + waitFor client.asyncClient.auth(token) + +proc ping*(client: SyncClient): bool = + waitFor client.asyncClient.ping() + proc `$`*(qr: QueryResult): string = if qr.columns.len == 0: return "(no results)" result = "" diff --git a/clients/python/baradb.py b/clients/python/baradb.py index 27edff4..8ff7a7a 100644 --- a/clients/python/baradb.py +++ b/clients/python/baradb.py @@ -15,6 +15,14 @@ Quick Start: for row in result: print(row["name"]) client.close() + +Parameterized Queries: + result = client.query_params("SELECT * FROM users WHERE age > $1", [WireValue.int64(18)]) + +Authentication: + client = Client("localhost", 9472, username="admin", password="secret") + client.connect() + client.auth("jwt-token-here") """ import socket @@ -41,17 +49,27 @@ class FieldKind: class MsgKind: + # Client messages + CLIENT_HANDSHAKE = 0x01 QUERY = 0x02 + QUERY_PARAMS = 0x03 + EXECUTE = 0x04 BATCH = 0x05 TRANSACTION = 0x06 CLOSE = 0x07 PING = 0x08 AUTH = 0x09 # Server messages + SERVER_HANDSHAKE = 0x80 READY = 0x81 DATA = 0x82 COMPLETE = 0x83 ERROR = 0x84 + AUTH_CHALLENGE = 0x85 + AUTH_OK = 0x86 + SCHEMA_CHANGE = 0x87 + PONG = 0x88 + TRANSACTION_STATE = 0x89 class ResultFormat: @@ -69,27 +87,105 @@ class WireValue: def null(): return WireValue(FieldKind.NULL) + @staticmethod + def bool_val(val: bool): + return WireValue(FieldKind.BOOL, val) + + @staticmethod + def int8(val: int): + return WireValue(FieldKind.INT8, val) + + @staticmethod + def int16(val: int): + return WireValue(FieldKind.INT16, val) + + @staticmethod + def int32(val: int): + return WireValue(FieldKind.INT32, val) + @staticmethod def int64(val: int): return WireValue(FieldKind.INT64, val) @staticmethod - def string(val: str): - return WireValue(FieldKind.STRING, val) + def float32(val: float): + return WireValue(FieldKind.FLOAT32, val) @staticmethod def float64(val: float): return WireValue(FieldKind.FLOAT64, val) @staticmethod - def bool_val(val: bool): - return WireValue(FieldKind.BOOL, val) + def string(val: str): + return WireValue(FieldKind.STRING, val) + + @staticmethod + def bytes_val(val: bytes): + return WireValue(FieldKind.BYTES, val) + + @staticmethod + def array_val(val: list): + return WireValue(FieldKind.ARRAY, val) + + @staticmethod + def object_val(val: dict): + return WireValue(FieldKind.OBJECT, val) + + @staticmethod + def vector(val: list): + return WireValue(FieldKind.VECTOR, val) + + @staticmethod + def json_val(val: str): + return WireValue(FieldKind.JSON, val) + + def serialize(self) -> bytes: + buf = bytes([self.kind]) + if self.kind == FieldKind.NULL: + pass + elif self.kind == FieldKind.BOOL: + buf += bytes([1 if self.value else 0]) + elif self.kind == FieldKind.INT8: + buf += struct.pack(">b", self.value) + elif self.kind == FieldKind.INT16: + buf += struct.pack(">h", self.value) + elif self.kind == FieldKind.INT32: + buf += struct.pack(">i", self.value) + elif self.kind == FieldKind.INT64: + buf += struct.pack(">q", self.value) + elif self.kind == FieldKind.FLOAT32: + buf += struct.pack(">f", self.value) + elif self.kind == FieldKind.FLOAT64: + buf += struct.pack(">d", self.value) + elif self.kind == FieldKind.STRING: + encoded = self.value.encode("utf-8") + buf += struct.pack(">I", len(encoded)) + encoded + elif self.kind == FieldKind.BYTES: + buf += struct.pack(">I", len(self.value)) + self.value + elif self.kind == FieldKind.ARRAY: + buf += struct.pack(">I", len(self.value)) + for item in self.value: + buf += item.serialize() + elif self.kind == FieldKind.OBJECT: + buf += struct.pack(">I", len(self.value)) + for key, val in self.value.items(): + key_bytes = key.encode("utf-8") + buf += struct.pack(">I", len(key_bytes)) + key_bytes + buf += val.serialize() + elif self.kind == FieldKind.VECTOR: + buf += struct.pack(">I", len(self.value)) + for f in self.value: + buf += struct.pack(">f", f) + elif self.kind == FieldKind.JSON: + encoded = self.value.encode("utf-8") + buf += struct.pack(">I", len(encoded)) + encoded + return buf class QueryResult: def __init__(self): self.columns: list[str] = [] - self.column_types: list[str] = [] + self.column_types: list[int] = [] self.rows: list[list[Any]] = [] self.row_count: int = 0 self.affected_rows: int = 0 @@ -130,6 +226,11 @@ class Client: def close(self) -> None: if self._sock: + try: + msg = self._build_message(MsgKind.CLOSE, b"") + self._sock.send(msg) + except Exception: + pass self._sock.close() self._connected = False @@ -140,6 +241,100 @@ class Client: self._request_id += 1 return self._request_id + def _recv_exact(self, size: int) -> bytes: + """Receive exactly `size` bytes from the socket.""" + data = b"" + while len(data) < size: + chunk = self._sock.recv(size - len(data)) + if not chunk: + raise ConnectionError("Connection closed by server") + data += chunk + return data + + def _read_response_header(self) -> tuple[int, int, int]: + """Read a 12-byte message header. Returns (kind, length, request_id).""" + header = self._recv_exact(12) + kind, length, req_id = struct.unpack(">III", header) + return kind, length, req_id + + def _read_error(self, length: int) -> Exception: + """Read and parse an ERROR payload.""" + data = self._recv_exact(length) + code = struct.unpack(">I", data[:4])[0] + msg_len = struct.unpack(">I", data[4:8])[0] + error_msg = data[8:8 + msg_len].decode("utf-8") + return Exception(f"BaraDB error {code}: {error_msg}") + + def _read_data_response(self, length: int) -> QueryResult: + """Read and parse a DATA payload, then follow up with COMPLETE.""" + data = self._recv_exact(length) + pos = [0] + + col_count = struct.unpack(">I", data[pos[0]:pos[0]+4])[0] + pos[0] += 4 + + cols = [] + for _ in range(col_count): + cols.append(self._read_string(data, pos)) + + col_types = [] + for _ in range(col_count): + col_types.append(data[pos[0]]) + pos[0] += 1 + + row_count = struct.unpack(">I", data[pos[0]:pos[0]+4])[0] + pos[0] += 4 + + rows = [] + for _ in range(row_count): + row = [] + for _ in range(col_count): + val = self._deserialize_value(data, pos) + row.append(val) + rows.append(row) + + result = QueryResult() + result.columns = cols + result.column_types = col_types + result.rows = rows + result.row_count = row_count + + comp_kind, comp_len, _ = self._read_response_header() + if comp_kind == MsgKind.COMPLETE: + comp_data = self._recv_exact(comp_len) + result.affected_rows = struct.unpack(">I", comp_data[:4])[0] + elif comp_kind == MsgKind.ERROR: + raise self._read_error(comp_len) + + return result + + def auth(self, token: str) -> None: + """Authenticate with the server using a JWT token.""" + encoded = token.encode("utf-8") + payload = struct.pack(">I", len(encoded)) + encoded + msg = self._build_message(MsgKind.AUTH, payload) + self._sock.send(msg) + + kind, length, _ = self._read_response_header() + if kind == MsgKind.AUTH_OK: + return + elif kind == MsgKind.ERROR: + raise self._read_error(length) + else: + raise Exception(f"Unexpected auth response: 0x{kind:02x}") + + def ping(self) -> bool: + """Ping the server. Returns True if pong received.""" + msg = self._build_message(MsgKind.PING, b"") + self._sock.send(msg) + + kind, length, _ = self._read_response_header() + if kind == MsgKind.PONG: + return True + elif kind == MsgKind.ERROR: + raise self._read_error(length) + return False + def query(self, sql: str) -> QueryResult: """Execute a BaraQL query.""" payload = self._encode_string(sql) @@ -148,58 +343,48 @@ class Client: msg = self._build_message(MsgKind.QUERY, payload) self._sock.send(msg) - result = QueryResult() - - # Read response header - header = self._sock.recv(12) - kind, length, req_id = struct.unpack(">III", header) + kind, length, _ = self._read_response_header() if kind == MsgKind.ERROR: - error_data = self._sock.recv(length) - code, msg_len = struct.unpack(">II", error_data[:8]) - error_msg = error_data[8:8 + msg_len].decode() - raise Exception(f"BaraDB error {code}: {error_msg}") + raise self._read_error(length) if kind == MsgKind.DATA: - data = self._sock.recv(length) - pos = [0] - col_count = struct.unpack(">I", data[pos[0]:pos[0]+4])[0] - pos[0] += 4 - cols = [] - for _ in range(col_count): - s = self._read_string(data, pos) - cols.append(s) - row_count = struct.unpack(">I", data[pos[0]:pos[0]+4])[0] - pos[0] += 4 - rows = [] - for _ in range(row_count): - row = [] - for _ in range(col_count): - val = self._deserialize_value(data, pos) - row.append(val) - rows.append(row) - col_types = [] - for _ in range(col_count): - col_types.append(hex(data[pos[0]])) - pos[0] += 1 - result.columns = cols - result.column_types = col_types - result.rows = rows - result.row_count = row_count - # Read following COMPLETE message - comp_header = self._sock.recv(12) - ckind, clen, _ = struct.unpack(">III", comp_header) - if ckind == MsgKind.COMPLETE: - comp_data = self._sock.recv(clen) - result.affected_rows = struct.unpack(">I", comp_data[:4])[0] - return result + return self._read_data_response(length) if kind == MsgKind.COMPLETE: - comp_data = self._sock.recv(length) - result.affected_rows = struct.unpack(">I", comp_data[:4])[0] + data = self._recv_exact(length) + result = QueryResult() + result.affected_rows = struct.unpack(">I", data[:4])[0] return result - return result + return QueryResult() + + def query_params(self, sql: str, params: Sequence[WireValue]) -> QueryResult: + """Execute a parameterized BaraQL query.""" + payload = self._encode_string(sql) + payload += bytes([ResultFormat.BINARY]) + payload += struct.pack(">I", len(params)) + for p in params: + payload += p.serialize() + + msg = self._build_message(MsgKind.QUERY_PARAMS, payload) + self._sock.send(msg) + + kind, length, _ = self._read_response_header() + + if kind == MsgKind.ERROR: + raise self._read_error(length) + + if kind == MsgKind.DATA: + return self._read_data_response(length) + + if kind == MsgKind.COMPLETE: + data = self._recv_exact(length) + result = QueryResult() + result.affected_rows = struct.unpack(">I", data[:4])[0] + return result + + return QueryResult() def execute(self, sql: str) -> int: result = self.query(sql) @@ -376,50 +561,3 @@ class QueryBuilder: def exec(self) -> QueryResult: return self.client.query(self.build()) - - -# BaraDB Binary Protocol Specification -""" -Protocol Format: - -Each message: [kind: uint32] [length: uint32] [request_id: uint32] [payload: bytes...] - -Client Messages: - 0x01 CLIENT_HANDSHAKE - 0x02 QUERY (string query, uint8 format) - 0x03 QUERY_PARAMS (string query, uint16 param_count, params...) - 0x04 EXECUTE (prepared statement) - 0x05 BATCH (batch of queries) - 0x06 TRANSACTION (begin/commit/rollback) - 0x07 CLOSE - 0x08 PING - 0x09 AUTH (auth method, credentials) - -Server Messages: - 0x80 SERVER_HANDSHAKE - 0x81 READY (transaction state) - 0x82 DATA (column count, column names, rows) - 0x83 COMPLETE (affected rows) - 0x84 ERROR (error code, error message) - 0x85 AUTH_CHALLENGE - 0x86 AUTH_OK - 0x87 SCHEMA_CHANGE - 0x88 PONG - 0x89 TRANSACTION_STATE - -Value Encoding: - value ::= kind:uint8 + data - NULL: 0x00 - BOOL: 0x01 + uint8(0|1) - INT8: 0x02 + int8 - INT16: 0x03 + int16(big-endian) - INT32: 0x04 + int32(big-endian) - INT64: 0x05 + int64(big-endian) - FLOAT32: 0x06 + float32(ieee754) - FLOAT64: 0x07 + float64(ieee754) - STRING: 0x08 + uint32(length) + utf8bytes - BYTES: 0x09 + uint32(length) + bytes - ARRAY: 0x0A + uint32(count) + value* - OBJECT: 0x0B + uint32(count) + (string_key + value)* - VECTOR: 0x0C + uint32(dim) + float32* -""" diff --git a/clients/rust/src/lib.rs b/clients/rust/src/lib.rs index f4c9210..7027654 100644 --- a/clients/rust/src/lib.rs +++ b/clients/rust/src/lib.rs @@ -1,31 +1,61 @@ //! BaraDB Rust Client //! //! Binary protocol client for BaraDB database. +//! Zero external dependencies — uses only `std`. //! //! # Example //! ```no_run -//! use baradb::Client; +//! use baradb::{Client, WireValue}; //! //! let mut client = Client::connect("localhost", 9472).unwrap(); //! let result = client.query("SELECT name FROM users WHERE age > 18").unwrap(); //! for row in result.rows() { -//! println!("{}", row["name"]); +//! if let Some(WireValue::String(name)) = row.get("name") { +//! println!("{}", name); +//! } //! } //! client.close(); //! ``` +//! +//! # Parameterized Queries +//! ```no_run +//! use baradb::{Client, WireValue}; +//! +//! let mut client = Client::connect("localhost", 9472).unwrap(); +//! let result = client.query_params( +//! "SELECT * FROM users WHERE age > $1", +//! &[WireValue::Int64(18)], +//! ).unwrap(); +//! ``` use std::collections::HashMap; use std::io::{Read, Write}; use std::net::TcpStream; -// Wire protocol constants +// Client message kinds +const MK_CLIENT_HANDSHAKE: u32 = 0x01; const MK_QUERY: u32 = 0x02; +const MK_QUERY_PARAMS: u32 = 0x03; +const MK_EXECUTE: u32 = 0x04; +const MK_BATCH: u32 = 0x05; +const MK_TRANSACTION: u32 = 0x06; +const MK_CLOSE: u32 = 0x07; +const MK_PING: u32 = 0x08; +const MK_AUTH: u32 = 0x09; + +// Server message kinds +const MK_SERVER_HANDSHAKE: u32 = 0x80; const MK_READY: u32 = 0x81; const MK_DATA: u32 = 0x82; const MK_COMPLETE: u32 = 0x83; const MK_ERROR: u32 = 0x84; -const MK_PING: u32 = 0x08; +const MK_AUTH_CHALLENGE: u32 = 0x85; +const MK_AUTH_OK: u32 = 0x86; +const MK_SCHEMA_CHANGE: u32 = 0x87; +const MK_PONG: u32 = 0x88; +const MK_TRANSACTION_STATE: u32 = 0x89; +// Field kinds const FK_NULL: u8 = 0x00; const FK_BOOL: u8 = 0x01; const FK_INT8: u8 = 0x02; @@ -41,6 +71,103 @@ const FK_OBJECT: u8 = 0x0B; const FK_VECTOR: u8 = 0x0C; const FK_JSON: u8 = 0x0D; +/// A typed wire value matching the BaraDB wire protocol. +#[derive(Debug, Clone)] +pub enum WireValue { + Null, + Bool(bool), + Int8(i8), + Int16(i16), + Int32(i32), + Int64(i64), + Float32(f32), + Float64(f64), + String(String), + Bytes(Vec), + Array(Vec), + Object(Vec<(String, WireValue)>), + Vector(Vec), + Json(String), +} + +impl WireValue { + pub fn to_string_lossy(&self) -> String { + match self { + WireValue::Null => String::new(), + WireValue::Bool(v) => v.to_string(), + WireValue::Int8(v) => v.to_string(), + WireValue::Int16(v) => v.to_string(), + WireValue::Int32(v) => v.to_string(), + WireValue::Int64(v) => v.to_string(), + WireValue::Float32(v) => v.to_string(), + WireValue::Float64(v) => v.to_string(), + WireValue::String(v) => v.clone(), + WireValue::Bytes(v) => format!("", v.len()), + WireValue::Array(v) => format!("", v.len()), + WireValue::Object(v) => format!("", v.len()), + WireValue::Vector(v) => format!("", v.len()), + WireValue::Json(v) => v.clone(), + } + } + + pub fn serialize(&self) -> Vec { + let mut buf = Vec::new(); + self.serialize_into(&mut buf); + buf + } + + fn serialize_into(&self, buf: &mut Vec) { + match self { + WireValue::Null => buf.push(FK_NULL), + WireValue::Bool(v) => { buf.push(FK_BOOL); buf.push(if *v { 1 } else { 0 }); } + WireValue::Int8(v) => { buf.push(FK_INT8); buf.push(*v as u8); } + WireValue::Int16(v) => { buf.push(FK_INT16); buf.extend_from_slice(&v.to_be_bytes()); } + WireValue::Int32(v) => { buf.push(FK_INT32); buf.extend_from_slice(&v.to_be_bytes()); } + WireValue::Int64(v) => { buf.push(FK_INT64); buf.extend_from_slice(&v.to_be_bytes()); } + WireValue::Float32(v) => { buf.push(FK_FLOAT32); buf.extend_from_slice(&v.to_be_bytes()); } + WireValue::Float64(v) => { buf.push(FK_FLOAT64); buf.extend_from_slice(&v.to_be_bytes()); } + WireValue::String(v) => { + buf.push(FK_STRING); + buf.extend_from_slice(&(v.len() as u32).to_be_bytes()); + buf.extend_from_slice(v.as_bytes()); + } + WireValue::Bytes(v) => { + buf.push(FK_BYTES); + buf.extend_from_slice(&(v.len() as u32).to_be_bytes()); + buf.extend_from_slice(v); + } + WireValue::Array(v) => { + buf.push(FK_ARRAY); + buf.extend_from_slice(&(v.len() as u32).to_be_bytes()); + for item in v { + item.serialize_into(buf); + } + } + WireValue::Object(v) => { + buf.push(FK_OBJECT); + buf.extend_from_slice(&(v.len() as u32).to_be_bytes()); + for (key, val) in v { + buf.extend_from_slice(&(key.len() as u32).to_be_bytes()); + buf.extend_from_slice(key.as_bytes()); + val.serialize_into(buf); + } + } + WireValue::Vector(v) => { + buf.push(FK_VECTOR); + buf.extend_from_slice(&(v.len() as u32).to_be_bytes()); + for f in v { + buf.extend_from_slice(&f.to_be_bytes()); + } + } + WireValue::Json(v) => { + buf.push(FK_JSON); + buf.extend_from_slice(&(v.len() as u32).to_be_bytes()); + buf.extend_from_slice(v.as_bytes()); + } + } + } +} + /// Connection configuration #[derive(Debug, Clone)] pub struct Config { @@ -67,8 +194,8 @@ impl Default for Config { #[derive(Debug)] pub struct QueryResult { columns: Vec, - column_types: Vec, - rows: Vec>, + column_types: Vec, + rows: Vec>, affected_rows: usize, } @@ -77,11 +204,11 @@ impl QueryResult { &self.columns } - pub fn column_types(&self) -> &[String] { + pub fn column_types(&self) -> &[u8] { &self.column_types } - pub fn rows(&self) -> &[HashMap] { + pub fn rows(&self) -> &[HashMap] { &self.rows } @@ -103,7 +230,6 @@ pub struct Client { } impl Client { - /// Connect to BaraDB server pub fn connect(host: &str, port: u16) -> Result> { let config = Config { host: host.to_string(), @@ -113,7 +239,6 @@ impl Client { Self::connect_with_config(config) } - /// Connect with custom configuration pub fn connect_with_config(config: Config) -> Result> { let addr = format!("{}:{}", config.host, config.port); let stream = TcpStream::connect(&addr)?; @@ -125,12 +250,13 @@ impl Client { }) } - /// Close the connection pub fn close(&mut self) { + if self.connected { + let _ = self.send_close(); + } self.connected = false; } - /// Check if connected pub fn is_connected(&self) -> bool { self.connected } @@ -140,93 +266,182 @@ impl Client { self.request_id } - /// Execute a BaraQL query + fn send_close(&mut self) -> Result<(), Box> { + let msg = build_message(MK_CLOSE, self.next_id(), &[]); + self.stream.write_all(&msg)?; + Ok(()) + } + + fn read_header(&mut self) -> Result<(u32, u32, u32), Box> { + let mut header = [0u8; 12]; + self.stream.read_exact(&mut header)?; + let kind = u32::from_be_bytes([header[0], header[1], header[2], header[3]]); + let length = u32::from_be_bytes([header[4], header[5], header[6], header[7]]); + let req_id = u32::from_be_bytes([header[8], header[9], header[10], header[11]]); + Ok((kind, length, req_id)) + } + + fn read_payload(&mut self, length: u32) -> Result, Box> { + let mut payload = vec![0u8; length as usize]; + if length > 0 { + self.stream.read_exact(&mut payload)?; + } + Ok(payload) + } + + fn read_error_message(payload: &[u8]) -> String { + if payload.len() >= 8 { + let code = u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]); + let msg_len = u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize; + if payload.len() >= 8 + msg_len { + let msg = String::from_utf8_lossy(&payload[8..8 + msg_len]); + return format!("BaraDB error {}: {}", code, msg); + } + } + "Query error".to_string() + } + + fn read_data_response(&mut self, payload: &[u8]) -> Result> { + let mut pos = 0usize; + let col_count = read_u32(payload, &mut pos) as usize; + + let mut columns = Vec::with_capacity(col_count); + for _ in 0..col_count { + columns.push(read_string(payload, &mut pos)); + } + + let mut col_types = Vec::with_capacity(col_count); + for _ in 0..col_count { + col_types.push(payload[pos]); + pos += 1; + } + + let row_count = read_u32(payload, &mut pos) as usize; + let mut rows = Vec::with_capacity(row_count); + for _ in 0..row_count { + let mut row = HashMap::new(); + for c in 0..col_count { + let val = read_wire_value(payload, &mut pos); + row.insert(columns[c].clone(), val); + } + rows.push(row); + } + + let mut affected = 0usize; + let (comp_kind, comp_len, _) = self.read_header()?; + if comp_kind == MK_COMPLETE { + let comp_payload = self.read_payload(comp_len)?; + if comp_payload.len() >= 4 { + affected = u32::from_be_bytes([comp_payload[0], comp_payload[1], comp_payload[2], comp_payload[3]]) as usize; + } + } + + Ok(QueryResult { columns, column_types: col_types, rows, affected_rows: affected }) + } + + pub fn auth(&mut self, token: &str) -> Result<(), Box> { + if !self.connected { + return Err("Not connected".into()); + } + let payload = encode_string(token); + let msg = build_message(MK_AUTH, self.next_id(), &payload); + self.stream.write_all(&msg)?; + + let (kind, length, _) = self.read_header()?; + match kind { + MK_AUTH_OK => Ok(()), + MK_ERROR => { + let p = self.read_payload(length)?; + Err(Self::read_error_message(&p).into()) + } + _ => Err(format!("Unexpected auth response: 0x{:02x}", kind).into()), + } + } + + pub fn ping(&mut self) -> Result> { + if !self.connected { + return Err("Not connected".into()); + } + let msg = build_message(MK_PING, self.next_id(), &[]); + self.stream.write_all(&msg)?; + + let (kind, length, _) = self.read_header()?; + match kind { + MK_PONG => Ok(true), + MK_ERROR => { + let p = self.read_payload(length)?; + Err(Self::read_error_message(&p).into()) + } + _ => Ok(false), + } + } + pub fn query(&mut self, sql: &str) -> Result> { if !self.connected { return Err("Not connected".into()); } - let payload = encode_string(sql); + let mut payload = encode_string(sql); + payload.push(0x00); // ResultFormat::BINARY + let msg = build_message(MK_QUERY, self.next_id(), &payload); self.stream.write_all(&msg)?; - // Read header - let mut header = [0u8; 12]; - self.stream.read_exact(&mut header)?; - - let kind = u32::from_be_bytes([header[0], header[1], header[2], header[3]]); - let length = u32::from_be_bytes([header[4], header[5], header[6], header[7]]) as usize; - - // Read payload - let mut payload = vec![0u8; length]; - if length > 0 { - self.stream.read_exact(&mut payload)?; - } + let (kind, length, _) = self.read_header()?; + let resp_payload = self.read_payload(length)?; match kind { MK_READY => Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: 0 }), - MK_DATA => { - let mut pos = 0usize; - let col_count = read_u32(&payload, &mut pos) as usize; - let mut columns = Vec::with_capacity(col_count); - for _ in 0..col_count { - columns.push(read_string(&payload, &mut pos)); - } - let mut col_types = Vec::with_capacity(col_count); - for _ in 0..col_count { - col_types.push(format!("{}", payload[pos])); - pos += 1; - } - let row_count = read_u32(&payload, &mut pos) as usize; - let mut rows = Vec::with_capacity(row_count); - for _ in 0..row_count { - let mut row = HashMap::new(); - for c in 0..col_count { - let val = read_value(&payload, &mut pos); - row.insert(columns[c].clone(), val); - } - rows.push(row); - } - // Read following COMPLETE message - let mut comp_header = [0u8; 12]; - self.stream.read_exact(&mut comp_header)?; - let comp_kind = u32::from_be_bytes([comp_header[0], comp_header[1], comp_header[2], comp_header[3]]); - let comp_len = u32::from_be_bytes([comp_header[4], comp_header[5], comp_header[6], comp_header[7]]) as usize; - let mut comp_payload = vec![0u8; comp_len]; - if comp_len > 0 { - self.stream.read_exact(&mut comp_payload)?; - } - let affected = if comp_kind == MK_COMPLETE && comp_payload.len() >= 4 { - u32::from_be_bytes([comp_payload[0], comp_payload[1], comp_payload[2], comp_payload[3]]) as usize - } else { 0 }; - Ok(QueryResult { columns, column_types: col_types, rows, affected_rows: affected }) - } + MK_DATA => self.read_data_response(&resp_payload), MK_COMPLETE => { - let affected = if payload.len() >= 4 { - u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize + let affected = if resp_payload.len() >= 4 { + u32::from_be_bytes([resp_payload[0], resp_payload[1], resp_payload[2], resp_payload[3]]) as usize } else { 0 }; Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: affected }) } - MK_ERROR => Err("Query error".into()), + MK_ERROR => Err(Self::read_error_message(&resp_payload).into()), + _ => Err(format!("Unknown response kind: {}", kind).into()), + } + } + + pub fn query_params(&mut self, sql: &str, params: &[WireValue]) -> Result> { + if !self.connected { + return Err("Not connected".into()); + } + + let mut payload = encode_string(sql); + payload.push(0x00); // ResultFormat::BINARY + payload.extend_from_slice(&(params.len() as u32).to_be_bytes()); + for p in params { + p.serialize_into(&mut payload); + } + + let msg = build_message(MK_QUERY_PARAMS, self.next_id(), &payload); + self.stream.write_all(&msg)?; + + let (kind, length, _) = self.read_header()?; + let resp_payload = self.read_payload(length)?; + + match kind { + MK_READY => Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: 0 }), + MK_DATA => self.read_data_response(&resp_payload), + MK_COMPLETE => { + let affected = if resp_payload.len() >= 4 { + u32::from_be_bytes([resp_payload[0], resp_payload[1], resp_payload[2], resp_payload[3]]) as usize + } else { 0 }; + Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: affected }) + } + MK_ERROR => Err(Self::read_error_message(&resp_payload).into()), _ => Err(format!("Unknown response kind: {}", kind).into()), } } - /// Execute a statement pub fn execute(&mut self, sql: &str) -> Result> { let result = self.query(sql)?; Ok(result.affected_rows()) } - - /// Ping the server - pub fn ping(&mut self) -> Result<(), Box> { - let msg = build_message(MK_PING, self.next_id(), &[]); - self.stream.write_all(&msg)?; - Ok(()) - } } -/// Query builder for fluent API pub struct QueryBuilder<'a> { client: &'a mut Client, select_cols: Vec, @@ -362,7 +577,6 @@ fn encode_string(s: &str) -> Vec { result } - fn read_u32(data: &[u8], pos: &mut usize) -> u32 { let val = u32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); *pos += 4; @@ -376,75 +590,88 @@ fn read_string(data: &[u8], pos: &mut usize) -> String { s } -fn read_value(data: &[u8], pos: &mut usize) -> String { +fn read_wire_value(data: &[u8], pos: &mut usize) -> WireValue { let kind = data[*pos]; *pos += 1; match kind { - FK_NULL => String::new(), + FK_NULL => WireValue::Null, FK_BOOL => { let val = data[*pos] != 0; *pos += 1; - val.to_string() + WireValue::Bool(val) } FK_INT8 => { let val = data[*pos] as i8; *pos += 1; - val.to_string() + WireValue::Int8(val) } FK_INT16 => { let val = i16::from_be_bytes([data[*pos], data[*pos + 1]]); *pos += 2; - val.to_string() + WireValue::Int16(val) } FK_INT32 => { let val = i32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); *pos += 4; - val.to_string() + WireValue::Int32(val) } FK_INT64 => { - let val = i64::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3], - data[*pos + 4], data[*pos + 5], data[*pos + 6], data[*pos + 7]]); + let val = i64::from_be_bytes([ + data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3], + data[*pos + 4], data[*pos + 5], data[*pos + 6], data[*pos + 7], + ]); *pos += 8; - val.to_string() + WireValue::Int64(val) } FK_FLOAT32 => { let val = f32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); *pos += 4; - val.to_string() + WireValue::Float32(val) } FK_FLOAT64 => { - let val = f64::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3], - data[*pos + 4], data[*pos + 5], data[*pos + 6], data[*pos + 7]]); + let val = f64::from_be_bytes([ + data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3], + data[*pos + 4], data[*pos + 5], data[*pos + 6], data[*pos + 7], + ]); *pos += 8; - val.to_string() + WireValue::Float64(val) } - FK_STRING => read_string(data, pos), + FK_STRING => WireValue::String(read_string(data, pos)), FK_BYTES => { let len = read_u32(data, pos) as usize; + let bytes = data[*pos..*pos + len].to_vec(); *pos += len; - format!("", len) + WireValue::Bytes(bytes) } FK_ARRAY => { - let count = read_u32(data, pos); + let count = read_u32(data, pos) as usize; + let mut arr = Vec::with_capacity(count); for _ in 0..count { - let _ = read_value(data, pos); + arr.push(read_wire_value(data, pos)); } - format!("", count) + WireValue::Array(arr) } FK_OBJECT => { - let count = read_u32(data, pos); + let count = read_u32(data, pos) as usize; + let mut obj = Vec::with_capacity(count); for _ in 0..count { - let _ = read_string(data, pos); - let _ = read_value(data, pos); + let key = read_string(data, pos); + let val = read_wire_value(data, pos); + obj.push((key, val)); } - format!("", count) + WireValue::Object(obj) } FK_VECTOR => { - let dim = read_u32(data, pos); - *pos += (dim * 4) as usize; - format!("", dim) + let dim = read_u32(data, pos) as usize; + let mut vec = Vec::with_capacity(dim); + for _ in 0..dim { + let f = f32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); + *pos += 4; + vec.push(f); + } + WireValue::Vector(vec) } - FK_JSON => read_string(data, pos), - _ => String::new(), + FK_JSON => WireValue::Json(read_string(data, pos)), + _ => WireValue::Null, } }