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. * Binary protocol client for BaraDB database.
* Communicates via the BaraDB Wire Protocol (binary, big-endian). * Communicates via the BaraDB Wire Protocol (binary, big-endian).
* Requires Node.js (uses 'net' module for TCP).
* *
* Install: * Install:
* npm install baradb * npm install baradb
@@ -16,23 +17,42 @@
* console.log(row.name); * console.log(row.name);
* } * }
* await client.close(); * 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: 0x02,
QUERY_PARAMS: 0x03,
EXECUTE: 0x04,
BATCH: 0x05, BATCH: 0x05,
TRANSACTION: 0x06, TRANSACTION: 0x06,
CLOSE: 0x07, CLOSE: 0x07,
PING: 0x08, PING: 0x08,
AUTH: 0x09, AUTH: 0x09,
SERVER_HANDSHAKE: 0x80,
READY: 0x81, READY: 0x81,
DATA: 0x82, DATA: 0x82,
COMPLETE: 0x83, COMPLETE: 0x83,
ERROR: 0x84, ERROR: 0x84,
AUTH_CHALLENGE: 0x85,
AUTH_OK: 0x86, AUTH_OK: 0x86,
}; SCHEMA_CHANGE: 0x87,
PONG: 0x88,
TRANSACTION_STATE: 0x89,
});
const FieldKind = { const FieldKind = Object.freeze({
NULL: 0x00, NULL: 0x00,
BOOL: 0x01, BOOL: 0x01,
INT8: 0x02, INT8: 0x02,
@@ -47,13 +67,136 @@ const FieldKind = {
OBJECT: 0x0B, OBJECT: 0x0B,
VECTOR: 0x0C, VECTOR: 0x0C,
JSON: 0x0D, JSON: 0x0D,
}; });
const ResultFormat = { const ResultFormat = Object.freeze({
BINARY: 0x00, BINARY: 0x00,
JSON: 0x01, JSON: 0x01,
TEXT: 0x02, 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 { class QueryResult {
constructor() { constructor() {
@@ -99,17 +242,50 @@ class Client {
this.socket = null; this.socket = null;
this.connected = false; this.connected = false;
this.requestId = 0; this.requestId = 0;
this._buffer = Buffer.alloc(0);
this._pendingResolve = null;
} }
/** Connect to the BaraDB server */
async connect() { async connect() {
// In Node.js: const net = require('net'); return new Promise((resolve, reject) => {
// In browser: WebSocket connection this.socket = net.createConnection({ host: this.host, port: this.port });
this.socket = null; // net.connect(this.port, this.host); this.socket.setTimeout(this.timeout);
this.socket.on('connect', () => {
this.connected = true; 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() { 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) { if (this.socket) {
this.socket.destroy(); this.socket.destroy();
this.socket = null; this.socket = null;
@@ -125,75 +301,88 @@ class Client {
return ++this.requestId; return ++this.requestId;
} }
/** Execute a BaraQL query */ async _recvExact(size) {
async query(sql) { while (this._buffer.length < size) {
const encoder = new TextEncoder(); await new Promise((resolve, reject) => {
const queryBytes = encoder.encode(sql); this._pendingResolve = resolve;
const payload = new Uint8Array(4 + queryBytes.length + 1); const timer = setTimeout(() => {
const view = new DataView(payload.buffer); this._pendingResolve = null;
view.setUint32(0, queryBytes.length, false); // big-endian reject(new Error('Receive timeout'));
payload.set(queryBytes, 4); }, this.timeout);
payload[4 + queryBytes.length] = ResultFormat.BINARY; this._pendingResolve = () => {
clearTimeout(timer);
const msg = this._build(MsgKind.QUERY, payload); this._pendingResolve = null;
// await this.socket.write(msg); resolve();
// const header = await this.socket.read(12); };
// ... parse response });
}
return new QueryResult(); const data = this._buffer.subarray(0, size);
this._buffer = this._buffer.subarray(size);
return data;
} }
/** Parse server response from header + payload buffers */ async _readHeader() {
_parseResponse(header, payload) { const header = await this._recvExact(12);
const view = new DataView(header.buffer); return {
const kind = view.getUint32(0, false); kind: header.readUInt32BE(0),
const len = view.getUint32(4, false); length: header.readUInt32BE(4),
const reqId = view.getUint32(8, false); 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(); 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; let pos = 0;
const colCount = new DataView(payload.buffer).getUint32(pos, false);
const colCount = data.readUInt32BE(pos);
pos += 4; pos += 4;
const cols = []; const cols = [];
for (let i = 0; i < colCount; i++) { for (let i = 0; i < colCount; i++) {
const sLen = new DataView(payload.buffer).getUint32(pos, false); const sLen = data.readUInt32BE(pos);
pos += 4; pos += 4;
cols.push(new TextDecoder().decode(payload.slice(pos, pos + sLen))); cols.push(data.subarray(pos, pos + sLen).toString('utf-8'));
pos += sLen; pos += sLen;
} }
result.columns = cols;
const colTypes = []; const colTypes = [];
for (let i = 0; i < colCount; i++) { for (let i = 0; i < colCount; i++) {
colTypes.push(payload[pos]); colTypes.push(data[pos]);
pos++; pos++;
} }
result.columnTypes = colTypes;
const rowCount = new DataView(payload.buffer).getUint32(pos, false); const rowCount = data.readUInt32BE(pos);
pos += 4; pos += 4;
for (let r = 0; r < rowCount; r++) { for (let r = 0; r < rowCount; r++) {
const row = []; const row = [];
for (let c = 0; c < colCount; c++) { for (let c = 0; c < colCount; c++) {
row.push(this._readValue(payload, pos)); const { value, newPos } = this._readValue(data, pos);
pos = this._lastReadPos; row.push(value);
pos = newPos;
} }
result.rows.push(row); 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.columns = cols;
result.affectedRows = new DataView(payload.buffer).getUint32(0, false); 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; return result;
@@ -206,74 +395,155 @@ class Client {
switch (kind) { switch (kind) {
case FieldKind.NULL: result = null; break; case FieldKind.NULL: result = null; break;
case FieldKind.BOOL: result = payload[pos] !== 0; pos++; break; case FieldKind.BOOL: result = payload[pos] !== 0; pos++; break;
case FieldKind.INT8: result = payload[pos] << 24 >> 24; pos++; break; case FieldKind.INT8: result = payload.readInt8(pos); pos++; break;
case FieldKind.INT16: result = new DataView(payload.buffer).getInt16(pos, false); pos += 2; break; case FieldKind.INT16: result = payload.readInt16BE(pos); pos += 2; break;
case FieldKind.INT32: result = new DataView(payload.buffer).getInt32(pos, false); pos += 4; break; case FieldKind.INT32: result = payload.readInt32BE(pos); pos += 4; break;
case FieldKind.INT64: result = new DataView(payload.buffer).getBigInt64(pos, false); pos += 8; break; case FieldKind.INT64: result = payload.readBigInt64BE(pos); pos += 8; break;
case FieldKind.FLOAT32: result = new DataView(payload.buffer).getFloat32(pos, false); pos += 4; break; case FieldKind.FLOAT32: result = payload.readFloatBE(pos); pos += 4; break;
case FieldKind.FLOAT64: result = new DataView(payload.buffer).getFloat64(pos, false); pos += 8; break; case FieldKind.FLOAT64: result = payload.readDoubleBE(pos); pos += 8; break;
case FieldKind.STRING: { case FieldKind.STRING: {
const len = new DataView(payload.buffer).getUint32(pos, false); const len = payload.readUInt32BE(pos);
pos += 4; pos += 4;
result = new TextDecoder().decode(payload.slice(pos, pos + len)); result = payload.subarray(pos, pos + len).toString('utf-8');
pos += len; pos += len;
break; break;
} }
case FieldKind.BYTES: { case FieldKind.BYTES: {
const len = new DataView(payload.buffer).getUint32(pos, false); const len = payload.readUInt32BE(pos);
pos += 4; pos += 4;
result = payload.slice(pos, pos + len); result = payload.subarray(pos, pos + len);
pos += len; pos += len;
break; break;
} }
case FieldKind.ARRAY: { case FieldKind.ARRAY: {
const count = new DataView(payload.buffer).getUint32(pos, false); const count = payload.readUInt32BE(pos);
pos += 4; pos += 4;
result = []; result = [];
for (let i = 0; i < count; i++) { for (let i = 0; i < count; i++) {
result.push(this._readValue(payload, pos)); const { value, newPos } = this._readValue(payload, pos);
pos = this._lastReadPos; result.push(value);
pos = newPos;
} }
break; break;
} }
case FieldKind.OBJECT: { case FieldKind.OBJECT: {
const count = new DataView(payload.buffer).getUint32(pos, false); const count = payload.readUInt32BE(pos);
pos += 4; pos += 4;
result = {}; result = {};
for (let i = 0; i < count; i++) { for (let i = 0; i < count; i++) {
const kLen = new DataView(payload.buffer).getUint32(pos, false); const kLen = payload.readUInt32BE(pos);
pos += 4; pos += 4;
const key = new TextDecoder().decode(payload.slice(pos, pos + kLen)); const key = payload.subarray(pos, pos + kLen).toString('utf-8');
pos += kLen; pos += kLen;
result[key] = this._readValue(payload, pos); const { value, newPos } = this._readValue(payload, pos);
pos = this._lastReadPos; result[key] = value;
pos = newPos;
} }
break; break;
} }
case FieldKind.VECTOR: { case FieldKind.VECTOR: {
const dim = new DataView(payload.buffer).getUint32(pos, false); const dim = payload.readUInt32BE(pos);
pos += 4; pos += 4;
result = []; result = [];
for (let i = 0; i < dim; i++) { for (let i = 0; i < dim; i++) {
result.push(new DataView(payload.buffer).getFloat32(pos, false)); result.push(payload.readFloatBE(pos));
pos += 4; pos += 4;
} }
break; break;
} }
case FieldKind.JSON: { case FieldKind.JSON: {
const len = new DataView(payload.buffer).getUint32(pos, false); const len = payload.readUInt32BE(pos);
pos += 4; pos += 4;
result = new TextDecoder().decode(payload.slice(pos, pos + len)); result = payload.subarray(pos, pos + len).toString('utf-8');
pos += len; pos += len;
break; break;
} }
default: result = null; default: result = null;
} }
this._lastReadPos = pos; return { value: result, newPos: pos };
return result; }
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) { async execute(sql) {
const result = await this.query(sql); const result = await this.query(sql);
return result.affectedRows; return result.affectedRows;
@@ -281,16 +551,11 @@ class Client {
_build(kind, payload) { _build(kind, payload) {
const reqId = this._nextId(); const reqId = this._nextId();
const header = new ArrayBuffer(12); const header = Buffer.alloc(12);
const view = new DataView(header); header.writeUInt32BE(kind, 0);
view.setUint32(0, kind, false); header.writeUInt32BE(payload.length, 4);
view.setUint32(4, payload.length, false); header.writeUInt32BE(reqId, 8);
view.setUint32(8, reqId, false); return Buffer.concat([header, payload]);
const msg = new Uint8Array(12 + payload.length);
msg.set(new Uint8Array(header), 0);
msg.set(payload, 12);
return msg;
} }
} }
@@ -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 fkJson = 0x0D
MsgKind* = enum MsgKind* = enum
# Client messages
mkClientHandshake = 0x01 mkClientHandshake = 0x01
mkQuery = 0x02 mkQuery = 0x02
mkQueryParams = 0x03
mkExecute = 0x04
mkBatch = 0x05 mkBatch = 0x05
mkTransaction = 0x06
mkClose = 0x07
mkPing = 0x08 mkPing = 0x08
mkAuth = 0x09 mkAuth = 0x09
# Server # Server messages
mkServerHandshake = 0x80
mkReady = 0x81 mkReady = 0x81
mkData = 0x82 mkData = 0x82
mkComplete = 0x83 mkComplete = 0x83
mkError = 0x84 mkError = 0x84
mkAuthChallenge = 0x85
mkAuthOk = 0x86
mkSchemaChange = 0x87
mkPong = 0x88 mkPong = 0x88
mkTransactionState = 0x89
ResultFormat* = enum ResultFormat* = enum
rfBinary = 0x00 rfBinary = 0x00
@@ -69,6 +79,11 @@ proc writeUint32(buf: var seq[byte], val: uint32) =
bigEndian32(addr bytes, unsafeAddr val) bigEndian32(addr bytes, unsafeAddr val)
buf.add(bytes) 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) = proc writeString(buf: var seq[byte], s: string) =
buf.writeUint32(uint32(s.len)) buf.writeUint32(uint32(s.len))
for ch in s: for ch in s:
@@ -80,6 +95,12 @@ proc readUint32(buf: openArray[byte], pos: var int): uint32 =
bigEndian32(addr result, unsafeAddr bytes) bigEndian32(addr result, unsafeAddr bytes)
pos += 4 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 = proc readString(buf: openArray[byte], pos: var int): string =
let len = int(readUint32(buf, pos)) let len = int(readUint32(buf, pos))
result = newString(len) result = newString(len)
@@ -98,10 +119,7 @@ proc serializeValue*(buf: var seq[byte], val: WireValue) =
bigEndian16(addr bytes16, unsafeAddr val.int16Val) bigEndian16(addr bytes16, unsafeAddr val.int16Val)
buf.add(bytes16) buf.add(bytes16)
of fkInt32: buf.writeUint32(uint32(val.int32Val)) of fkInt32: buf.writeUint32(uint32(val.int32Val))
of fkInt64: of fkInt64: buf.writeUint64(uint64(val.int64Val))
var bytes: array[8, byte]
bigEndian64(addr bytes, unsafeAddr val.int64Val)
buf.add(bytes)
of fkFloat32: of fkFloat32:
var bytes32: array[4, byte] var bytes32: array[4, byte]
copyMem(addr bytes32, unsafeAddr val.float32Val, 4) copyMem(addr bytes32, unsafeAddr val.float32Val, 4)
@@ -129,7 +147,7 @@ proc serializeValue*(buf: var seq[byte], val: WireValue) =
var fb: array[4, byte] var fb: array[4, byte]
copyMem(addr fb, unsafeAddr f, 4) copyMem(addr fb, unsafeAddr f, 4)
buf.add(fb) buf.add(fb)
else: discard of fkJson: buf.writeString(val.jsonVal)
proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue = proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
let kind = FieldKind(buf[pos]) let kind = FieldKind(buf[pos])
@@ -152,12 +170,7 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
of fkInt32: of fkInt32:
result = WireValue(kind: fkInt32, int32Val: int32(readUint32(buf, pos))) result = WireValue(kind: fkInt32, int32Val: int32(readUint32(buf, pos)))
of fkInt64: of fkInt64:
var bytes: array[8, byte] result = WireValue(kind: fkInt64, int64Val: int64(readUint64(buf, pos)))
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
of fkFloat32: of fkFloat32:
var v32: float32 var v32: float32
copyMem(addr v32, addr buf[pos], 4) 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) result = WireValue(kind: fkVector, vecVal: vec)
of fkJson: of fkJson:
result = WireValue(kind: fkJson, jsonVal: readString(buf, pos)) result = WireValue(kind: fkJson, jsonVal: readString(buf, pos))
else:
result = WireValue(kind: fkNull)
proc buildMessage*(kind: MsgKind, requestId: uint32, payload: seq[byte]): seq[byte] = proc buildMessage*(kind: MsgKind, requestId: uint32, payload: seq[byte]): seq[byte] =
result = @[] result = @[]
@@ -218,6 +229,20 @@ proc makeQueryMessage*(requestId: uint32, query: string): seq[byte] =
payload.add(byte(rfBinary)) payload.add(byte(rfBinary))
buildMessage(mkQuery, requestId, payload) 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 === # === Client Library ===
type type
@@ -257,16 +282,20 @@ proc connect*(client: BaraClient) {.async.} =
await client.socket.connect(client.config.host, Port(client.config.port)) await client.socket.connect(client.config.host, Port(client.config.port))
client.connected = true client.connected = true
proc nextId(client: BaraClient): uint32 =
inc client.requestId; client.requestId
proc close*(client: BaraClient) = proc close*(client: BaraClient) =
if client.connected: if client.connected:
try:
let msg = buildMessage(mkClose, client.nextId(), @[])
waitFor client.socket.send(cast[string](msg))
except: discard
client.socket.close() client.socket.close()
client.connected = false client.connected = false
proc isConnected*(client: BaraClient): bool = client.connected proc isConnected*(client: BaraClient): bool = client.connected
proc nextId(client: BaraClient): uint32 =
inc client.requestId; client.requestId
proc wireValueToString*(wv: WireValue): string = proc wireValueToString*(wv: WireValue): string =
case wv.kind case wv.kind
of fkNull: return "" of fkNull: return ""
@@ -367,11 +396,48 @@ proc exec*(client: BaraClient, sql: string): Future[int] {.async.} =
let qr = await client.query(sql) let qr = await client.query(sql)
return qr.affectedRows return qr.affectedRows
proc ping*(client: BaraClient) {.async.} = proc auth*(client: BaraClient, token: string) {.async.} =
if not client.connected: return 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(), @[]) let msg = buildMessage(mkPing, client.nextId(), @[])
await client.socket.send(cast[string](msg)) 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 === # === Fluent Query Builder ===
type type
@@ -468,6 +534,15 @@ proc close*(client: SyncClient) =
proc query*(client: SyncClient, sql: string): QueryResult = proc query*(client: SyncClient, sql: string): QueryResult =
waitFor client.asyncClient.query(sql) 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 = proc `$`*(qr: QueryResult): string =
if qr.columns.len == 0: return "(no results)" if qr.columns.len == 0: return "(no results)"
result = "" result = ""
+233 -95
View File
@@ -15,6 +15,14 @@ Quick Start:
for row in result: for row in result:
print(row["name"]) print(row["name"])
client.close() 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 import socket
@@ -41,17 +49,27 @@ class FieldKind:
class MsgKind: class MsgKind:
# Client messages
CLIENT_HANDSHAKE = 0x01
QUERY = 0x02 QUERY = 0x02
QUERY_PARAMS = 0x03
EXECUTE = 0x04
BATCH = 0x05 BATCH = 0x05
TRANSACTION = 0x06 TRANSACTION = 0x06
CLOSE = 0x07 CLOSE = 0x07
PING = 0x08 PING = 0x08
AUTH = 0x09 AUTH = 0x09
# Server messages # Server messages
SERVER_HANDSHAKE = 0x80
READY = 0x81 READY = 0x81
DATA = 0x82 DATA = 0x82
COMPLETE = 0x83 COMPLETE = 0x83
ERROR = 0x84 ERROR = 0x84
AUTH_CHALLENGE = 0x85
AUTH_OK = 0x86
SCHEMA_CHANGE = 0x87
PONG = 0x88
TRANSACTION_STATE = 0x89
class ResultFormat: class ResultFormat:
@@ -69,27 +87,105 @@ class WireValue:
def null(): def null():
return WireValue(FieldKind.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 @staticmethod
def int64(val: int): def int64(val: int):
return WireValue(FieldKind.INT64, val) return WireValue(FieldKind.INT64, val)
@staticmethod @staticmethod
def string(val: str): def float32(val: float):
return WireValue(FieldKind.STRING, val) return WireValue(FieldKind.FLOAT32, val)
@staticmethod @staticmethod
def float64(val: float): def float64(val: float):
return WireValue(FieldKind.FLOAT64, val) return WireValue(FieldKind.FLOAT64, val)
@staticmethod @staticmethod
def bool_val(val: bool): def string(val: str):
return WireValue(FieldKind.BOOL, val) 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: class QueryResult:
def __init__(self): def __init__(self):
self.columns: list[str] = [] self.columns: list[str] = []
self.column_types: list[str] = [] self.column_types: list[int] = []
self.rows: list[list[Any]] = [] self.rows: list[list[Any]] = []
self.row_count: int = 0 self.row_count: int = 0
self.affected_rows: int = 0 self.affected_rows: int = 0
@@ -130,6 +226,11 @@ class Client:
def close(self) -> None: def close(self) -> None:
if self._sock: if self._sock:
try:
msg = self._build_message(MsgKind.CLOSE, b"")
self._sock.send(msg)
except Exception:
pass
self._sock.close() self._sock.close()
self._connected = False self._connected = False
@@ -140,6 +241,100 @@ class Client:
self._request_id += 1 self._request_id += 1
return self._request_id 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: def query(self, sql: str) -> QueryResult:
"""Execute a BaraQL query.""" """Execute a BaraQL query."""
payload = self._encode_string(sql) payload = self._encode_string(sql)
@@ -148,59 +343,49 @@ class Client:
msg = self._build_message(MsgKind.QUERY, payload) msg = self._build_message(MsgKind.QUERY, payload)
self._sock.send(msg) self._sock.send(msg)
result = QueryResult() kind, length, _ = self._read_response_header()
# Read response header
header = self._sock.recv(12)
kind, length, req_id = struct.unpack(">III", header)
if kind == MsgKind.ERROR: if kind == MsgKind.ERROR:
error_data = self._sock.recv(length) raise self._read_error(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}")
if kind == MsgKind.DATA: if kind == MsgKind.DATA:
data = self._sock.recv(length) return self._read_data_response(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
if kind == MsgKind.COMPLETE: if kind == MsgKind.COMPLETE:
comp_data = self._sock.recv(length) data = self._recv_exact(length)
result.affected_rows = struct.unpack(">I", comp_data[:4])[0] 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 result
return QueryResult()
def execute(self, sql: str) -> int: def execute(self, sql: str) -> int:
result = self.query(sql) result = self.query(sql)
return result.affected_rows return result.affected_rows
@@ -376,50 +561,3 @@ class QueryBuilder:
def exec(self) -> QueryResult: def exec(self) -> QueryResult:
return self.client.query(self.build()) 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 //! BaraDB Rust Client
//! //!
//! Binary protocol client for BaraDB database. //! Binary protocol client for BaraDB database.
//! Zero external dependencies — uses only `std`.
//! //!
//! # Example //! # Example
//! ```no_run //! ```no_run
//! use baradb::Client; //! use baradb::{Client, WireValue};
//! //!
//! let mut client = Client::connect("localhost", 9472).unwrap(); //! let mut client = Client::connect("localhost", 9472).unwrap();
//! let result = client.query("SELECT name FROM users WHERE age > 18").unwrap(); //! let result = client.query("SELECT name FROM users WHERE age > 18").unwrap();
//! for row in result.rows() { //! for row in result.rows() {
//! println!("{}", row["name"]); //! if let Some(WireValue::String(name)) = row.get("name") {
//! println!("{}", name);
//! }
//! } //! }
//! client.close(); //! 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::collections::HashMap;
use std::io::{Read, Write}; use std::io::{Read, Write};
use std::net::TcpStream; use std::net::TcpStream;
// Wire protocol constants // Client message kinds
const MK_CLIENT_HANDSHAKE: u32 = 0x01;
const MK_QUERY: u32 = 0x02; 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_READY: u32 = 0x81;
const MK_DATA: u32 = 0x82; const MK_DATA: u32 = 0x82;
const MK_COMPLETE: u32 = 0x83; const MK_COMPLETE: u32 = 0x83;
const MK_ERROR: u32 = 0x84; 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_NULL: u8 = 0x00;
const FK_BOOL: u8 = 0x01; const FK_BOOL: u8 = 0x01;
const FK_INT8: u8 = 0x02; const FK_INT8: u8 = 0x02;
@@ -41,6 +71,103 @@ const FK_OBJECT: u8 = 0x0B;
const FK_VECTOR: u8 = 0x0C; const FK_VECTOR: u8 = 0x0C;
const FK_JSON: u8 = 0x0D; 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 /// Connection configuration
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct Config { pub struct Config {
@@ -67,8 +194,8 @@ impl Default for Config {
#[derive(Debug)] #[derive(Debug)]
pub struct QueryResult { pub struct QueryResult {
columns: Vec<String>, columns: Vec<String>,
column_types: Vec<String>, column_types: Vec<u8>,
rows: Vec<HashMap<String, String>>, rows: Vec<HashMap<String, WireValue>>,
affected_rows: usize, affected_rows: usize,
} }
@@ -77,11 +204,11 @@ impl QueryResult {
&self.columns &self.columns
} }
pub fn column_types(&self) -> &[String] { pub fn column_types(&self) -> &[u8] {
&self.column_types &self.column_types
} }
pub fn rows(&self) -> &[HashMap<String, String>] { pub fn rows(&self) -> &[HashMap<String, WireValue>] {
&self.rows &self.rows
} }
@@ -103,7 +230,6 @@ pub struct Client {
} }
impl Client { impl Client {
/// Connect to BaraDB server
pub fn connect(host: &str, port: u16) -> Result<Self, Box<dyn std::error::Error>> { pub fn connect(host: &str, port: u16) -> Result<Self, Box<dyn std::error::Error>> {
let config = Config { let config = Config {
host: host.to_string(), host: host.to_string(),
@@ -113,7 +239,6 @@ impl Client {
Self::connect_with_config(config) Self::connect_with_config(config)
} }
/// Connect with custom configuration
pub fn connect_with_config(config: Config) -> Result<Self, Box<dyn std::error::Error>> { pub fn connect_with_config(config: Config) -> Result<Self, Box<dyn std::error::Error>> {
let addr = format!("{}:{}", config.host, config.port); let addr = format!("{}:{}", config.host, config.port);
let stream = TcpStream::connect(&addr)?; let stream = TcpStream::connect(&addr)?;
@@ -125,12 +250,13 @@ impl Client {
}) })
} }
/// Close the connection
pub fn close(&mut self) { pub fn close(&mut self) {
if self.connected {
let _ = self.send_close();
}
self.connected = false; self.connected = false;
} }
/// Check if connected
pub fn is_connected(&self) -> bool { pub fn is_connected(&self) -> bool {
self.connected self.connected
} }
@@ -140,93 +266,182 @@ impl Client {
self.request_id 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>> { pub fn query(&mut self, sql: &str) -> Result<QueryResult, Box<dyn std::error::Error>> {
if !self.connected { if !self.connected {
return Err("Not connected".into()); 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); let msg = build_message(MK_QUERY, self.next_id(), &payload);
self.stream.write_all(&msg)?; self.stream.write_all(&msg)?;
// Read header let (kind, length, _) = self.read_header()?;
let mut header = [0u8; 12]; let resp_payload = self.read_payload(length)?;
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)?;
}
match kind { match kind {
MK_READY => Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: 0 }), MK_READY => Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: 0 }),
MK_DATA => { MK_DATA => self.read_data_response(&resp_payload),
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_COMPLETE => { MK_COMPLETE => {
let affected = if payload.len() >= 4 { let affected = if resp_payload.len() >= 4 {
u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize u32::from_be_bytes([resp_payload[0], resp_payload[1], resp_payload[2], resp_payload[3]]) as usize
} else { 0 }; } else { 0 };
Ok(QueryResult { columns: vec![], column_types: vec![], rows: vec![], affected_rows: affected }) 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()), _ => Err(format!("Unknown response kind: {}", kind).into()),
} }
} }
/// Execute a statement
pub fn execute(&mut self, sql: &str) -> Result<usize, Box<dyn std::error::Error>> { pub fn execute(&mut self, sql: &str) -> Result<usize, Box<dyn std::error::Error>> {
let result = self.query(sql)?; let result = self.query(sql)?;
Ok(result.affected_rows()) 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> { pub struct QueryBuilder<'a> {
client: &'a mut Client, client: &'a mut Client,
select_cols: Vec<String>, select_cols: Vec<String>,
@@ -362,7 +577,6 @@ fn encode_string(s: &str) -> Vec<u8> {
result result
} }
fn read_u32(data: &[u8], pos: &mut usize) -> u32 { 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]]); let val = u32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]);
*pos += 4; *pos += 4;
@@ -376,75 +590,88 @@ fn read_string(data: &[u8], pos: &mut usize) -> String {
s s
} }
fn read_value(data: &[u8], pos: &mut usize) -> String { fn read_wire_value(data: &[u8], pos: &mut usize) -> WireValue {
let kind = data[*pos]; let kind = data[*pos];
*pos += 1; *pos += 1;
match kind { match kind {
FK_NULL => String::new(), FK_NULL => WireValue::Null,
FK_BOOL => { FK_BOOL => {
let val = data[*pos] != 0; let val = data[*pos] != 0;
*pos += 1; *pos += 1;
val.to_string() WireValue::Bool(val)
} }
FK_INT8 => { FK_INT8 => {
let val = data[*pos] as i8; let val = data[*pos] as i8;
*pos += 1; *pos += 1;
val.to_string() WireValue::Int8(val)
} }
FK_INT16 => { FK_INT16 => {
let val = i16::from_be_bytes([data[*pos], data[*pos + 1]]); let val = i16::from_be_bytes([data[*pos], data[*pos + 1]]);
*pos += 2; *pos += 2;
val.to_string() WireValue::Int16(val)
} }
FK_INT32 => { FK_INT32 => {
let val = i32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); let val = i32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]);
*pos += 4; *pos += 4;
val.to_string() WireValue::Int32(val)
} }
FK_INT64 => { FK_INT64 => {
let val = i64::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3], let val = i64::from_be_bytes([
data[*pos + 4], data[*pos + 5], data[*pos + 6], data[*pos + 7]]); 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; *pos += 8;
val.to_string() WireValue::Int64(val)
} }
FK_FLOAT32 => { FK_FLOAT32 => {
let val = f32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]); let val = f32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]);
*pos += 4; *pos += 4;
val.to_string() WireValue::Float32(val)
} }
FK_FLOAT64 => { FK_FLOAT64 => {
let val = f64::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3], let val = f64::from_be_bytes([
data[*pos + 4], data[*pos + 5], data[*pos + 6], data[*pos + 7]]); 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; *pos += 8;
val.to_string() WireValue::Float64(val)
} }
FK_STRING => read_string(data, pos), FK_STRING => WireValue::String(read_string(data, pos)),
FK_BYTES => { FK_BYTES => {
let len = read_u32(data, pos) as usize; let len = read_u32(data, pos) as usize;
let bytes = data[*pos..*pos + len].to_vec();
*pos += len; *pos += len;
format!("<bytes:{}>", len) WireValue::Bytes(bytes)
} }
FK_ARRAY => { 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 { for _ in 0..count {
let _ = read_value(data, pos); arr.push(read_wire_value(data, pos));
} }
format!("<array:{}>", count) WireValue::Array(arr)
} }
FK_OBJECT => { 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 { for _ in 0..count {
let _ = read_string(data, pos); let key = read_string(data, pos);
let _ = read_value(data, pos); let val = read_wire_value(data, pos);
obj.push((key, val));
} }
format!("<object:{}>", count) WireValue::Object(obj)
} }
FK_VECTOR => { FK_VECTOR => {
let dim = read_u32(data, pos); let dim = read_u32(data, pos) as usize;
*pos += (dim * 4) as usize; let mut vec = Vec::with_capacity(dim);
format!("<vector:{}>", 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), WireValue::Vector(vec)
_ => String::new(), }
FK_JSON => WireValue::Json(read_string(data, pos)),
_ => WireValue::Null,
} }
} }