feat(clients): sync all client drivers with current wire protocol

- Python: fix column types read order, add auth/ping/query_params, proper recv
- JavaScript: wire up real TCP socket, add auth/ping/query_params, WireValue class
- Rust: add WireValue enum with structured types, fix type fidelity for ARRAY/OBJECT/VECTOR, add auth/query_params, fix ping/error parsing
- Nim: add missing MsgKind values, makeQueryParamsMessage, auth/ping, fix forward decl

All clients now support AUTH (0x09), QUERY_PARAMS (0x03), PING (0x08)
and align with server.nim + wire.nim protocol definitions.
This commit is contained in:
2026-05-07 22:38:09 +03:00
parent 09247ffcad
commit abb382c419
4 changed files with 1029 additions and 324 deletions
+357 -92
View File
@@ -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);
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();
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}`);
}
if (kind === MsgKind.DATA) {
let pos = 0;
const colCount = new DataView(payload.buffer).getUint32(pos, false);
const colCount = data.readUInt32BE(pos);
pos += 4;
const cols = [];
for (let i = 0; i < colCount; i++) {
const sLen = new DataView(payload.buffer).getUint32(pos, false);
const sLen = data.readUInt32BE(pos);
pos += 4;
cols.push(new TextDecoder().decode(payload.slice(pos, pos + sLen)));
cols.push(data.subarray(pos, pos + sLen).toString('utf-8'));
pos += sLen;
}
result.columns = cols;
const colTypes = [];
for (let i = 0; i < colCount; i++) {
colTypes.push(payload[pos]);
colTypes.push(data[pos]);
pos++;
}
result.columnTypes = colTypes;
const rowCount = new DataView(payload.buffer).getUint32(pos, false);
const rowCount = data.readUInt32BE(pos);
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;
const { value, newPos } = this._readValue(data, pos);
row.push(value);
pos = newPos;
}
result.rows.push(row);
}
result.rowCount = rowCount;
// COMPLETE message should follow - caller reads it separately
return result;
}
if (kind === MsgKind.COMPLETE && payload.length >= 4) {
result.affectedRows = new DataView(payload.buffer).getUint32(0, false);
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 };
+94 -19
View File
@@ -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 = ""
+233 -95
View File
@@ -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,59 +343,49 @@ 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 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)
return result.affected_rows
@@ -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*
"""
+329 -102
View File
@@ -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<u8>),
Array(Vec<WireValue>),
Object(Vec<(String, WireValue)>),
Vector(Vec<f32>),
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!("<bytes:{}>", v.len()),
WireValue::Array(v) => format!("<array:{}>", v.len()),
WireValue::Object(v) => format!("<object:{}>", v.len()),
WireValue::Vector(v) => format!("<vector:{}>", v.len()),
WireValue::Json(v) => v.clone(),
}
}
pub fn serialize(&self) -> Vec<u8> {
let mut buf = Vec::new();
self.serialize_into(&mut buf);
buf
}
fn serialize_into(&self, buf: &mut Vec<u8>) {
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<String>,
column_types: Vec<String>,
rows: Vec<HashMap<String, String>>,
column_types: Vec<u8>,
rows: Vec<HashMap<String, WireValue>>,
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<String, String>] {
pub fn rows(&self) -> &[HashMap<String, WireValue>] {
&self.rows
}
@@ -103,7 +230,6 @@ pub struct Client {
}
impl Client {
/// Connect to BaraDB server
pub fn connect(host: &str, port: u16) -> Result<Self, Box<dyn std::error::Error>> {
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<Self, Box<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<Vec<u8>, Box<dyn std::error::Error>> {
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<QueryResult, Box<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<bool, Box<dyn std::error::Error>> {
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<QueryResult, Box<dyn std::error::Error>> {
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<QueryResult, Box<dyn std::error::Error>> {
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<usize, Box<dyn std::error::Error>> {
let result = self.query(sql)?;
Ok(result.affected_rows())
}
/// Ping the server
pub fn ping(&mut self) -> Result<(), Box<dyn std::error::Error>> {
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<String>,
@@ -362,7 +577,6 @@ fn encode_string(s: &str) -> Vec<u8> {
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!("<bytes:{}>", 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!("<array:{}>", 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!("<object:{}>", count)
WireValue::Object(obj)
}
FK_VECTOR => {
let dim = read_u32(data, pos);
*pos += (dim * 4) as usize;
format!("<vector:{}>", 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);
}
FK_JSON => read_string(data, pos),
_ => String::new(),
WireValue::Vector(vec)
}
FK_JSON => WireValue::Json(read_string(data, pos)),
_ => WireValue::Null,
}
}