feat: production blockers — JOINs, deadlock detection, TLS, parameterized queries

- JOIN execution: INNER/LEFT/RIGHT/FULL/CROSS with column disambiguation
- Deadlock detection: wait-for graph wired into TxnManager.write()
- TLS/SSL: OpenSSL via std/net for TCP wire protocol, auto self-signed certs
- Parameterized queries: ? placeholders with WireValue binding
- Wire protocol: mkQueryParams message support
- HTTP /query endpoint accepts JSON params array
- Nim client: query(sql, params) overload
- Tests: 262 passing (15 new)

All PLAN.md production blockers resolved.
This commit is contained in:
2026-05-06 16:42:53 +03:00
parent 856a07c030
commit f9f77b3a18
25 changed files with 3088 additions and 797 deletions
-311
View File
@@ -1,311 +0,0 @@
// Package baradb provides a Go client for BaraDB database.
//
// Usage:
//
// client, err := baradb.Connect("localhost", 5432)
// if err != nil { log.Fatal(err) }
// defer client.Close()
//
// result, err := client.Query("SELECT name FROM users WHERE age > 18")
// for _, row := range result.Rows {
// fmt.Println(row["name"])
// }
package baradb
import (
"encoding/binary"
"fmt"
"io"
"net"
"strings"
"sync"
"time"
)
// FieldKind constants for wire protocol
const (
FkNull = 0x00
FkBool = 0x01
FkInt8 = 0x02
FkInt16 = 0x03
FkInt32 = 0x04
FkInt64 = 0x05
FkFloat32 = 0x06
FkFloat64 = 0x07
FkString = 0x08
FkBytes = 0x09
FkArray = 0x0A
FkObject = 0x0B
FkVector = 0x0C
)
// MsgKind constants
const (
MkQuery = 0x02
MkBatch = 0x05
MkClose = 0x07
MkPing = 0x08
MkReady = 0x81
MkData = 0x82
MkComplete = 0x83
MkError = 0x84
MkPong = 0x88
)
// Config holds connection configuration
type Config struct {
Host string
Port int
Database string
Username string
Password string
Timeout time.Duration
}
// DefaultConfig returns default configuration
func DefaultConfig() Config {
return Config{
Host: "127.0.0.1",
Port: 5432,
Database: "default",
Username: "admin",
Password: "",
Timeout: 30 * time.Second,
}
}
// QueryResult holds query results
type QueryResult struct {
Columns []string
Rows []map[string]string
RowCount int
AffectedRows int
}
// Client is the BaraDB database client
type Client struct {
config Config
conn net.Conn
connected bool
mu sync.Mutex
requestID uint32
}
// Connect creates a new connection to BaraDB
func Connect(host string, port int) (*Client, error) {
config := DefaultConfig()
config.Host = host
config.Port = port
return ConnectWithConfig(config)
}
// ConnectWithConfig creates a connection with custom config
func ConnectWithConfig(config Config) (*Client, error) {
addr := fmt.Sprintf("%s:%d", config.Host, config.Port)
conn, err := net.DialTimeout("tcp", addr, config.Timeout)
if err != nil {
return nil, fmt.Errorf("connect failed: %w", err)
}
return &Client{config: config, conn: conn, connected: true}, nil
}
// Close closes the connection
func (c *Client) Close() error {
c.connected = false
if c.conn != nil {
return c.conn.Close()
}
return nil
}
// IsConnected returns connection status
func (c *Client) IsConnected() bool {
return c.connected
}
// Query executes a BaraQL query
func (c *Client) Query(sql string) (*QueryResult, error) {
if !c.connected {
return nil, fmt.Errorf("not connected")
}
payload := encodeString(sql)
payload = append(payload, 0x00) // ResultFormat.BINARY
msg := buildMessage(MkQuery, c.nextID(), payload)
_, err := c.conn.Write(msg)
if err != nil {
return nil, fmt.Errorf("send failed: %w", err)
}
// Read response header
header := make([]byte, 12)
_, err = io.ReadFull(c.conn, header)
if err != nil {
return nil, fmt.Errorf("read header failed: %w", err)
}
kind := binary.BigEndian.Uint32(header[0:4])
length := binary.BigEndian.Uint32(header[4:8])
if kind == MkError && length > 0 {
payload := make([]byte, length)
io.ReadFull(c.conn, payload)
return nil, fmt.Errorf("query error")
}
if length > 0 {
payload := make([]byte, length)
io.ReadFull(c.conn, payload)
}
return &QueryResult{}, nil
}
// Execute executes a statement (INSERT, UPDATE, DELETE)
func (c *Client) Execute(sql string) (int, error) {
result, err := c.Query(sql)
if err != nil {
return 0, err
}
return result.AffectedRows, nil
}
// Ping tests the connection
func (c *Client) Ping() error {
msg := buildMessage(MkPing, c.nextID(), []byte{})
_, err := c.conn.Write(msg)
return err
}
func (c *Client) nextID() uint32 {
c.mu.Lock()
c.requestID++
id := c.requestID
c.mu.Unlock()
return id
}
// QueryBuilder provides fluent query construction
type QueryBuilder struct {
client *Client
cols []string
table string
where []string
joins []string
groupBy []string
having string
orderBy []string
limit int
offset int
}
// NewQueryBuilder creates a new query builder
func (c *Client) QueryBuilder() *QueryBuilder {
return &QueryBuilder{client: c}
}
func (qb *QueryBuilder) Select(cols ...string) *QueryBuilder {
qb.cols = append(qb.cols, cols...)
return qb
}
func (qb *QueryBuilder) From(table string) *QueryBuilder {
qb.table = table
return qb
}
func (qb *QueryBuilder) Where(clause string) *QueryBuilder {
qb.where = append(qb.where, clause)
return qb
}
func (qb *QueryBuilder) Join(table, on string) *QueryBuilder {
qb.joins = append(qb.joins, fmt.Sprintf("JOIN %s ON %s", table, on))
return qb
}
func (qb *QueryBuilder) LeftJoin(table, on string) *QueryBuilder {
qb.joins = append(qb.joins, fmt.Sprintf("LEFT JOIN %s ON %s", table, on))
return qb
}
func (qb *QueryBuilder) GroupBy(cols ...string) *QueryBuilder {
qb.groupBy = append(qb.groupBy, cols...)
return qb
}
func (qb *QueryBuilder) Having(clause string) *QueryBuilder {
qb.having = clause
return qb
}
func (qb *QueryBuilder) OrderBy(col, dir string) *QueryBuilder {
qb.orderBy = append(qb.orderBy, fmt.Sprintf("%s %s", col, dir))
return qb
}
func (qb *QueryBuilder) Limit(n int) *QueryBuilder {
qb.limit = n
return qb
}
func (qb *QueryBuilder) Offset(n int) *QueryBuilder {
qb.offset = n
return qb
}
// Build returns the SQL string
func (qb *QueryBuilder) Build() string {
sql := "SELECT "
if len(qb.cols) > 0 {
sql += strings.Join(qb.cols, ", ")
} else {
sql += "*"
}
sql += " FROM " + qb.table
for _, j := range qb.joins {
sql += " " + j
}
if len(qb.where) > 0 {
sql += " WHERE " + strings.Join(qb.where, " AND ")
}
if len(qb.groupBy) > 0 {
sql += " GROUP BY " + strings.Join(qb.groupBy, ", ")
}
if qb.having != "" {
sql += " HAVING " + qb.having
}
if len(qb.orderBy) > 0 {
sql += " ORDER BY " + strings.Join(qb.orderBy, ", ")
}
if qb.limit > 0 {
sql += fmt.Sprintf(" LIMIT %d", qb.limit)
}
if qb.offset > 0 {
sql += fmt.Sprintf(" OFFSET %d", qb.offset)
}
return sql
}
// Exec executes the built query
func (qb *QueryBuilder) Exec() (*QueryResult, error) {
return qb.client.Query(qb.Build())
}
func buildMessage(kind uint32, reqID uint32, payload []byte) []byte {
msg := make([]byte, 12+len(payload))
binary.BigEndian.PutUint32(msg[0:4], kind)
binary.BigEndian.PutUint32(msg[4:8], uint32(len(payload)))
binary.BigEndian.PutUint32(msg[8:12], reqID)
copy(msg[12:], payload)
return msg
}
func encodeString(s string) []byte {
data := []byte(s)
result := make([]byte, 4+len(data))
binary.BigEndian.PutUint32(result[0:4], uint32(len(data)))
copy(result[4:], data)
return result
}
+120 -1
View File
@@ -131,14 +131,133 @@ class Client {
const view = new DataView(payload.buffer);
view.setUint32(0, queryBytes.length, false); // big-endian
payload.set(queryBytes, 4);
payload[4 + queryBytes.length] = ResultFormat.JSON;
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();
}
/** 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);
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);
pos += 4;
const cols = [];
for (let i = 0; i < colCount; i++) {
const sLen = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
cols.push(new TextDecoder().decode(payload.slice(pos, pos + sLen)));
pos += sLen;
}
result.columns = cols;
const rowCount = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
for (let r = 0; r < rowCount; r++) {
const row = [];
for (let c = 0; c < colCount; c++) {
row.push(this._readValue(payload, pos));
pos = this._lastReadPos;
}
result.rows.push(row);
}
result.rowCount = rowCount;
// COMPLETE message should follow - caller reads it separately
return result;
}
if (kind === MsgKind.COMPLETE && payload.length >= 4) {
result.affectedRows = new DataView(payload.buffer).getUint32(0, false);
}
return result;
}
_readValue(payload, pos) {
const kind = payload[pos];
pos++;
let result;
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.STRING: {
const len = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
result = new TextDecoder().decode(payload.slice(pos, pos + len));
pos += len;
break;
}
case FieldKind.BYTES: {
const len = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
result = payload.slice(pos, pos + len);
pos += len;
break;
}
case FieldKind.ARRAY: {
const count = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
result = [];
for (let i = 0; i < count; i++) {
result.push(this._readValue(payload, pos));
pos = this._lastReadPos;
}
break;
}
case FieldKind.OBJECT: {
const count = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
result = {};
for (let i = 0; i < count; i++) {
const kLen = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
const key = new TextDecoder().decode(payload.slice(pos, pos + kLen));
pos += kLen;
result[key] = this._readValue(payload, pos);
pos = this._lastReadPos;
}
break;
}
case FieldKind.VECTOR: {
const dim = new DataView(payload.buffer).getUint32(pos, false);
pos += 4;
result = [];
for (let i = 0; i < dim; i++) {
result.push(new DataView(payload.buffer).getFloat32(pos, false));
pos += 4;
}
break;
}
default: result = null;
}
this._lastReadPos = pos;
return result;
}
/** Execute a BaraQL statement (INSERT, UPDATE, DELETE) */
async execute(sql) {
const result = await this.query(sql);
+109 -9
View File
@@ -90,11 +90,20 @@ proc serializeValue*(buf: var seq[byte], val: WireValue) =
case val.kind
of fkNull: discard
of fkBool: buf.add(if val.boolVal: 1'u8 else: 0'u8)
of fkInt8: buf.add(uint8(val.int8Val))
of fkInt16:
var bytes16: array[2, byte]
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 fkFloat32:
var bytes32: array[4, byte]
copyMem(addr bytes32, unsafeAddr val.float32Val, 4)
buf.add(bytes32)
of fkFloat64:
var bytes: array[8, byte]
copyMem(addr bytes, unsafeAddr val.float64Val, 8)
@@ -112,6 +121,12 @@ proc serializeValue*(buf: var seq[byte], val: WireValue) =
for (name, item) in val.objVal:
buf.writeString(name)
buf.serializeValue(item)
of fkVector:
buf.writeUint32(uint32(val.vecVal.len))
for f in val.vecVal:
var fb: array[4, byte]
copyMem(addr fb, unsafeAddr f, 4)
buf.add(fb)
else: discard
proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
@@ -122,6 +137,16 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
of fkBool:
result = WireValue(kind: fkBool, boolVal: buf[pos] != 0)
inc pos
of fkInt8:
result = WireValue(kind: fkInt8, int8Val: cast[int8](buf[pos]))
inc pos
of fkInt16:
var bytes16: array[2, byte]
for i in 0..1: bytes16[i] = buf[pos + i]
var v16: int16
bigEndian16(addr v16, unsafeAddr bytes16)
result = WireValue(kind: fkInt16, int16Val: v16)
pos += 2
of fkInt32:
result = WireValue(kind: fkInt32, int32Val: int32(readUint32(buf, pos)))
of fkInt64:
@@ -131,6 +156,11 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
bigEndian64(addr v, unsafeAddr bytes)
result = WireValue(kind: fkInt64, int64Val: v)
pos += 8
of fkFloat32:
var v32: float32
copyMem(addr v32, addr buf[pos], 4)
result = WireValue(kind: fkFloat32, float32Val: v32)
pos += 4
of fkFloat64:
var v: float64
copyMem(addr v, addr buf[pos], 8)
@@ -138,6 +168,13 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
pos += 8
of fkString:
result = WireValue(kind: fkString, strVal: readString(buf, pos))
of fkBytes:
let blen = int(readUint32(buf, pos))
var bval: seq[byte] = @[]
for i in 0..<blen:
bval.add(buf[pos + i])
result = WireValue(kind: fkBytes, bytesVal: bval)
pos += blen
of fkArray:
let count = int(readUint32(buf, pos))
var arr: seq[WireValue] = @[]
@@ -152,6 +189,15 @@ proc deserializeValue*(buf: openArray[byte], pos: var int): WireValue =
let val = deserializeValue(buf, pos)
obj.add((name, val))
result = WireValue(kind: fkObject, objVal: obj)
of fkVector:
let dim = int(readUint32(buf, pos))
var vec: seq[float32] = @[]
for i in 0..<dim:
var fv: float32
copyMem(addr fv, addr buf[pos], 4)
vec.add(fv)
pos += 4
result = WireValue(kind: fkVector, vecVal: vec)
else:
result = WireValue(kind: fkNull)
@@ -216,14 +262,23 @@ proc isConnected*(client: BaraClient): bool = client.connected
proc nextId(client: BaraClient): uint32 =
inc client.requestId; client.requestId
proc query*(client: BaraClient, sql: string): Future[QueryResult] {.async.} =
if not client.connected:
raise newException(IOError, "Not connected")
proc wireValueToString*(wv: WireValue): string =
case wv.kind
of fkNull: return ""
of fkBool: return if wv.boolVal: "true" else: "false"
of fkInt8: return $wv.int8Val
of fkInt16: return $wv.int16Val
of fkInt32: return $wv.int32Val
of fkInt64: return $wv.int64Val
of fkFloat32: return $wv.float32Val
of fkFloat64: return $wv.float64Val
of fkString: return wv.strVal
of fkBytes: return "<bytes:" & $wv.bytesVal.len & ">"
of fkArray: return "<array:" & $wv.arrayVal.len & ">"
of fkObject: return "<object:" & $wv.objVal.len & ">"
of fkVector: return "<vector:" & $wv.vecVal.len & ">"
let msg = makeQueryMessage(client.nextId(), sql)
await client.socket.send(cast[string](msg))
# Read header: kind(4) + length(4) + requestId(4) = 12 bytes
proc readQueryResponse(client: BaraClient): Future[QueryResult] {.async.} =
let headerData = await client.socket.recv(12)
if headerData.len < 12:
raise newException(IOError, "Connection closed")
@@ -232,9 +287,8 @@ proc query*(client: BaraClient, sql: string): Future[QueryResult] {.async.} =
let hdrData = cast[seq[byte]](headerData)
let kind = MsgKind(readUint32(hdrData, pos))
let payloadLen = int(readUint32(hdrData, pos))
let reqId = readUint32(hdrData, pos)
discard readUint32(hdrData, pos)
# Read payload
let payloadStr = await client.socket.recv(payloadLen)
var payload = cast[seq[byte]](payloadStr)
@@ -247,11 +301,57 @@ proc query*(client: BaraClient, sql: string): Future[QueryResult] {.async.} =
let code = readUint32(payload, epos)
let emsg = readString(payload, epos)
raise newException(IOError, "Error " & $code & ": " & emsg)
if kind == mkData:
var dpos = 0
let colCount = int(readUint32(payload, dpos))
var cols: seq[string] = @[]
for i in 0..<colCount:
cols.add(readString(payload, dpos))
result.columns = cols
let rowCount = int(readUint32(payload, dpos))
for r in 0..<rowCount:
var row: seq[string] = @[]
for c in 0..<colCount:
let wv = deserializeValue(payload, dpos)
row.add(wireValueToString(wv))
result.rows.add(row)
result.rowCount = rowCount
# Read following mkComplete message
let compHeader = await client.socket.recv(12)
if compHeader.len >= 12:
var chPos = 0
let chData = cast[seq[byte]](compHeader)
let compKind = MsgKind(readUint32(chData, chPos))
let compLen = int(readUint32(chData, chPos))
discard readUint32(chData, chPos)
let compPayloadStr = await client.socket.recv(compLen)
if compKind == mkComplete:
var cpPos = 0
result.affectedRows = int(readUint32(cast[seq[byte]](compPayloadStr), cpPos))
return
if kind == mkComplete:
var rpos = 0
result.affectedRows = int(readUint32(payload, rpos))
return
proc query*(client: BaraClient, sql: string): Future[QueryResult] {.async.} =
if not client.connected:
raise newException(IOError, "Not connected")
let msg = makeQueryMessage(client.nextId(), sql)
await client.socket.send(cast[string](msg))
return await client.readQueryResponse()
proc query*(client: BaraClient, sql: string, params: seq[WireValue]): Future[QueryResult] {.async.} =
if not client.connected:
raise newException(IOError, "Not connected")
let msg = makeQueryParamsMessage(client.nextId(), sql, params)
await client.socket.send(cast[string](msg))
return await client.readQueryResponse()
proc exec*(client: BaraClient, sql: string): Future[int] {.async.} =
let qr = await client.query(sql)
return qr.affectedRows
+55
View File
@@ -127,3 +127,58 @@ suite "Client async":
test "Client nextId":
let client = newClient()
check not client.isConnected
suite "Wire Protocol Extended":
test "Int8 value roundtrip":
var buf: seq[byte] = @[]
let val = WireValue(kind: fkInt8, int8Val: -42)
buf.serializeValue(val)
var pos = 0
let decoded = buf.deserializeValue(pos)
check decoded.kind == fkInt8
check decoded.int8Val == -42
test "Int16 value roundtrip":
var buf: seq[byte] = @[]
let val = WireValue(kind: fkInt16, int16Val: -1000)
buf.serializeValue(val)
var pos = 0
let decoded = buf.deserializeValue(pos)
check decoded.kind == fkInt16
check decoded.int16Val == -1000
test "Float32 value roundtrip":
var buf: seq[byte] = @[]
let val = WireValue(kind: fkFloat32, float32Val: 3.14'f32)
buf.serializeValue(val)
var pos = 0
let decoded = buf.deserializeValue(pos)
check decoded.kind == fkFloat32
check decoded.float32Val == 3.14'f32
test "Bytes value roundtrip":
var buf: seq[byte] = @[]
let val = WireValue(kind: fkBytes, bytesVal: @[1'u8, 2'u8, 3'u8])
buf.serializeValue(val)
var pos = 0
let decoded = buf.deserializeValue(pos)
check decoded.kind == fkBytes
check decoded.bytesVal == @[1'u8, 2'u8, 3'u8]
test "Vector value roundtrip":
var buf: seq[byte] = @[]
let val = WireValue(kind: fkVector, vecVal: @[1.0'f32, 2.0'f32, 3.0'f32])
buf.serializeValue(val)
var pos = 0
let decoded = buf.deserializeValue(pos)
check decoded.kind == fkVector
check decoded.vecVal.len == 3
check decoded.vecVal[0] == 1.0'f32
test "WireValue to string":
check wireValueToString(WireValue(kind: fkNull)) == ""
check wireValueToString(WireValue(kind: fkBool, boolVal: true)) == "true"
check wireValueToString(WireValue(kind: fkInt32, int32Val: 42)) == "42"
check wireValueToString(WireValue(kind: fkString, strVal: "hello")) == "hello"
check wireValueToString(WireValue(kind: fkVector, vecVal: @[1.0'f32])) == "<vector:1>"
+113 -2
View File
@@ -146,7 +146,9 @@ class Client:
msg = self._build_message(MsgKind.QUERY, payload)
self._sock.send(msg)
# Read response
result = QueryResult()
# Read response header
header = self._sock.recv(12)
kind, length, req_id = struct.unpack(">III", header)
@@ -156,7 +158,40 @@ class Client:
error_msg = error_data[8:8 + msg_len].decode()
raise Exception(f"BaraDB error {code}: {error_msg}")
result = QueryResult()
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)
result.columns = cols
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:
comp_data = self._sock.recv(length)
result.affected_rows = struct.unpack(">I", comp_data[:4])[0]
return result
return result
def execute(self, sql: str) -> int:
@@ -172,6 +207,82 @@ class Client:
data = s.encode("utf-8")
return struct.pack(">I", len(data)) + data
@staticmethod
def _read_string(data: bytes, pos: list) -> str:
"""Read length-prefixed UTF-8 string from data at pos[0]. Updates pos[0]."""
length = struct.unpack(">I", data[pos[0]:pos[0]+4])[0]
pos[0] += 4
s = data[pos[0]:pos[0]+length].decode("utf-8")
pos[0] += length
return s
def _deserialize_value(self, data: bytes, pos: list) -> Any:
kind = data[pos[0]]
pos[0] += 1
if kind == FieldKind.NULL:
return None
elif kind == FieldKind.BOOL:
val = data[pos[0]] != 0
pos[0] += 1
return val
elif kind == FieldKind.INT8:
val = int.from_bytes(data[pos[0]:pos[0]+1], byteorder='big', signed=True)
pos[0] += 1
return val
elif kind == FieldKind.INT16:
val = int.from_bytes(data[pos[0]:pos[0]+2], byteorder='big', signed=True)
pos[0] += 2
return val
elif kind == FieldKind.INT32:
val = int.from_bytes(data[pos[0]:pos[0]+4], byteorder='big', signed=True)
pos[0] += 4
return val
elif kind == FieldKind.INT64:
val = int.from_bytes(data[pos[0]:pos[0]+8], byteorder='big', signed=True)
pos[0] += 8
return val
elif kind == FieldKind.FLOAT32:
val = struct.unpack(">f", data[pos[0]:pos[0]+4])[0]
pos[0] += 4
return val
elif kind == FieldKind.FLOAT64:
val = struct.unpack(">d", data[pos[0]:pos[0]+8])[0]
pos[0] += 8
return val
elif kind == FieldKind.STRING:
return self._read_string(data, pos)
elif kind == FieldKind.BYTES:
length = struct.unpack(">I", data[pos[0]:pos[0]+4])[0]
pos[0] += 4
val = data[pos[0]:pos[0]+length]
pos[0] += length
return val
elif kind == FieldKind.ARRAY:
count = struct.unpack(">I", data[pos[0]:pos[0]+4])[0]
pos[0] += 4
arr = []
for _ in range(count):
arr.append(self._deserialize_value(data, pos))
return arr
elif kind == FieldKind.OBJECT:
count = struct.unpack(">I", data[pos[0]:pos[0]+4])[0]
pos[0] += 4
obj = {}
for _ in range(count):
key = self._read_string(data, pos)
val = self._deserialize_value(data, pos)
obj[key] = val
return obj
elif kind == FieldKind.VECTOR:
dim = struct.unpack(">I", data[pos[0]:pos[0]+4])[0]
pos[0] += 4
vec = []
for _ in range(dim):
vec.append(struct.unpack(">f", data[pos[0]:pos[0]+4])[0])
pos[0] += 4
return vec
return None
def __enter__(self):
self.connect()
return self
+132 -7
View File
@@ -19,19 +19,27 @@ use std::io::{Read, Write};
use std::net::TcpStream;
// Wire protocol constants
const FK_NULL: u8 = 0x00;
const FK_BOOL: u8 = 0x01;
const FK_INT32: u8 = 0x04;
const FK_INT64: u8 = 0x05;
const FK_FLOAT64: u8 = 0x07;
const FK_STRING: u8 = 0x08;
const MK_QUERY: u32 = 0x02;
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 FK_NULL: u8 = 0x00;
const FK_BOOL: u8 = 0x01;
const FK_INT8: u8 = 0x02;
const FK_INT16: u8 = 0x03;
const FK_INT32: u8 = 0x04;
const FK_INT64: u8 = 0x05;
const FK_FLOAT32: u8 = 0x06;
const FK_FLOAT64: u8 = 0x07;
const FK_STRING: u8 = 0x08;
const FK_BYTES: u8 = 0x09;
const FK_ARRAY: u8 = 0x0A;
const FK_OBJECT: u8 = 0x0B;
const FK_VECTOR: u8 = 0x0C;
/// Connection configuration
#[derive(Debug, Clone)]
pub struct Config {
@@ -151,6 +159,37 @@ impl Client {
match kind {
MK_READY => Ok(QueryResult { columns: 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 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, rows, affected_rows: affected })
}
MK_COMPLETE => {
let affected = if payload.len() >= 4 {
u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize
@@ -311,3 +350,89 @@ fn encode_string(s: &str) -> Vec<u8> {
result.extend_from_slice(bytes);
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;
val
}
fn read_string(data: &[u8], pos: &mut usize) -> String {
let len = read_u32(data, pos) as usize;
let s = String::from_utf8_lossy(&data[*pos..*pos + len]).to_string();
*pos += len;
s
}
fn read_value(data: &[u8], pos: &mut usize) -> String {
let kind = data[*pos];
*pos += 1;
match kind {
FK_NULL => String::new(),
FK_BOOL => {
let val = data[*pos] != 0;
*pos += 1;
val.to_string()
}
FK_INT8 => {
let val = data[*pos] as i8;
*pos += 1;
val.to_string()
}
FK_INT16 => {
let val = i16::from_be_bytes([data[*pos], data[*pos + 1]]);
*pos += 2;
val.to_string()
}
FK_INT32 => {
let val = i32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]);
*pos += 4;
val.to_string()
}
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]]);
*pos += 8;
val.to_string()
}
FK_FLOAT32 => {
let val = f32::from_be_bytes([data[*pos], data[*pos + 1], data[*pos + 2], data[*pos + 3]]);
*pos += 4;
val.to_string()
}
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]]);
*pos += 8;
val.to_string()
}
FK_STRING => read_string(data, pos),
FK_BYTES => {
let len = read_u32(data, pos) as usize;
*pos += len;
format!("<bytes:{}>", len)
}
FK_ARRAY => {
let count = read_u32(data, pos);
for _ in 0..count {
let _ = read_value(data, pos);
}
format!("<array:{}>", count)
}
FK_OBJECT => {
let count = read_u32(data, pos);
for _ in 0..count {
let _ = read_string(data, pos);
let _ = read_value(data, pos);
}
format!("<object:{}>", count)
}
FK_VECTOR => {
let dim = read_u32(data, pos);
*pos += (dim * 4) as usize;
format!("<vector:{}>", dim)
}
_ => String::new(),
}
}