fix(security): SQL injection in Python/JS clients + bare except + session leak
CI / test (push) Has been cancelled
CI / verify (push) Has been cancelled
Clients CI / build-server (push) Has been cancelled
Clients CI / test-python (push) Has been cancelled
Clients CI / test-javascript (push) Has been cancelled
Clients CI / test-nim (push) Has been cancelled
Clients CI / test-rust (push) Has been cancelled

- Fix SQL injection vulnerabilities in chat_history.py, langchain_store.py,
  rag_pipeline.py, baradb_langchain.js by switching to parameterized queries
- Replace dangerous bare except: with except CatchableError:/ValueError:
  in llm.nim, embed.nim, cypher.nim, mcp/server.nim, executor.nim
- Fix session variable leak in MCP handleVectorSearch/handleSchemaInspect
- Build: 0 errors, all tests pass
This commit is contained in:
2026-05-17 16:58:05 +03:00
parent c95bc4cd44
commit 1e38e29f25
9 changed files with 227 additions and 109 deletions
+47 -24
View File
@@ -2,7 +2,7 @@
* BaraDB LangChain.js Vector Store Integration * BaraDB LangChain.js Vector Store Integration
* *
* Usage: * Usage:
* const { Client } = require('./baradb'); * const { Client, WireValue } = require('./baradb');
* const { BaraDBStore } = require('./baradb_langchain'); * const { BaraDBStore } = require('./baradb_langchain');
* *
* const client = new Client('localhost', 9472); * const client = new Client('localhost', 9472);
@@ -24,6 +24,8 @@
* const results = await store.similaritySearch('hello', 5); * const results = await store.similaritySearch('hello', 5);
*/ */
const { WireValue } = require('./baradb');
class BaraDBStore { class BaraDBStore {
constructor(options = {}) { constructor(options = {}) {
this.client = options.client; this.client = options.client;
@@ -63,28 +65,23 @@ class BaraDBStore {
const vec = await this.embeddingFunction(text); const vec = await this.embeddingFunction(text);
const vecStr = '[' + vec.join(',') + ']'; const vecStr = '[' + vec.join(',') + ']';
const metaCols = []; const colNames = [this.embeddingCol, this.textCol];
const metaVals = []; const params = [WireValue.string(vecStr), WireValue.string(text)];
if (this.tenantId) { if (this.tenantId) {
metaCols.push('tenant_id'); colNames.push('tenant_id');
metaVals.push(`'${this.tenantId}'`); params.push(WireValue.string(this.tenantId));
} }
for (const mc of this.metadataCols) { for (const mc of this.metadataCols) {
if (meta[mc] !== undefined) { if (meta[mc] !== undefined) {
metaCols.push(mc); colNames.push(mc);
metaVals.push(`'${String(meta[mc]).replace(/'/g, "''")}'`); params.push(WireValue.string(String(meta[mc])));
} }
} }
let colList = `${this.embeddingCol}, ${this.textCol}`; const placeholders = params.map((_, i) => `$${i + 1}`).join(', ');
let valList = `'${vecStr}', '${text.replace(/'/g, "''")}'`; const sql = `INSERT INTO ${this.table} (${colNames.join(', ')}) VALUES (${placeholders}) RETURNING id`;
if (metaCols.length > 0) { const result = await this.client.queryParams(sql, params);
colList += ', ' + metaCols.join(', ');
valList += ', ' + metaVals.join(', ');
}
const sql = `INSERT INTO ${this.table} (${colList}) VALUES (${valList}) RETURNING id`;
const result = await this.client.query(sql);
if (result.rows && result.rows.length > 0) { if (result.rows && result.rows.length > 0) {
insertedIds.push(result.rows[0].id || result.rows[0][0]); insertedIds.push(result.rows[0].id || result.rows[0][0]);
} }
@@ -110,17 +107,36 @@ class BaraDBStore {
const vecStr = '[' + vec.join(',') + ']'; const vecStr = '[' + vec.join(',') + ']';
if (this.tenantId) { if (this.tenantId) {
await this.client.query(`SET app.tenant_id = '${this.tenantId}'`); await this.client.queryParams('SET app.tenant_id = $1', [WireValue.string(this.tenantId)]);
} }
let sql; let sql;
let params;
if (filter && filter.column && filter.value) { if (filter && filter.column && filter.value) {
sql = `SELECT hybrid_search_filtered('${this.table}', '${this.embeddingCol}', '${this.textCol}', '${query.replace(/'/g, "''")}', '${vecStr}', ${k}, '${filter.column}', '${filter.value}') AS res`; sql = 'SELECT hybrid_search_filtered($1, $2, $3, $4, $5, $6, $7, $8) AS res';
params = [
WireValue.string(this.table),
WireValue.string(this.embeddingCol),
WireValue.string(this.textCol),
WireValue.string(query),
WireValue.string(vecStr),
WireValue.int32(k),
WireValue.string(filter.column),
WireValue.string(filter.value),
];
} else { } else {
sql = `SELECT hybrid_search('${this.table}', '${this.embeddingCol}', '${this.textCol}', '${query.replace(/'/g, "''")}', '${vecStr}', ${k}) AS res`; sql = 'SELECT hybrid_search($1, $2, $3, $4, $5, $6) AS res';
params = [
WireValue.string(this.table),
WireValue.string(this.embeddingCol),
WireValue.string(this.textCol),
WireValue.string(query),
WireValue.string(vecStr),
WireValue.int32(k),
];
} }
const result = await this.client.query(sql); const result = await this.client.queryParams(sql, params);
if (!result.rows || result.rows.length === 0) return []; if (!result.rows || result.rows.length === 0) return [];
const raw = result.rows[0].res || result.rows[0][0] || '[]'; const raw = result.rows[0].res || result.rows[0][0] || '[]';
@@ -135,7 +151,10 @@ class BaraDBStore {
for (const item of arr) { for (const item of arr) {
const docId = item.id; const docId = item.id;
const score = parseFloat(item.score || 0); const score = parseFloat(item.score || 0);
const rowResult = await this.client.query(`SELECT * FROM ${this.table} WHERE id = ${docId}`); const rowResult = await this.client.queryParams(
`SELECT * FROM ${this.table} WHERE id = $1`,
[WireValue.string(String(docId))]
);
if (rowResult.rows && rowResult.rows.length > 0) { if (rowResult.rows && rowResult.rows.length > 0) {
const row = rowResult.rows[0]; const row = rowResult.rows[0];
const pageContent = row[this.textCol] || row[Object.keys(row).find(k => k.toLowerCase() === this.textCol.toLowerCase())]; const pageContent = row[this.textCol] || row[Object.keys(row).find(k => k.toLowerCase() === this.textCol.toLowerCase())];
@@ -180,13 +199,17 @@ class BaraDBStore {
async delete(ids) { async delete(ids) {
await this._ensureTable(); await this._ensureTable();
if (!ids || ids.length === 0) return; if (!ids || ids.length === 0) return;
const idList = ids.join(', '); const params = ids.map((id, i) => WireValue.string(String(id)));
await this.client.query(`DELETE FROM ${this.table} WHERE id IN (${idList})`); const placeholders = ids.map((_, i) => `$${i + 1}`).join(', ');
await this.client.queryParams(
`DELETE FROM ${this.table} WHERE id IN (${placeholders})`,
params
);
} }
async setTenant(tenantId) { async setTenant(tenantId) {
this.tenantId = tenantId; this.tenantId = tenantId;
await this.client.query(`SET app.tenant_id = '${tenantId}'`); await this.client.queryParams('SET app.tenant_id = $1', [WireValue.string(tenantId)]);
} }
} }
+59 -26
View File
@@ -5,7 +5,7 @@ Implements LangChain's BaseChatMessageHistory interface backed by BaraDB.
Supports multi-tenant isolation via tenant_id and user_id. Supports multi-tenant isolation via tenant_id and user_id.
Usage: Usage:
from baradb import Client from baradb import Client, WireValue
from baradb.chat_history import BaraDBChatHistory from baradb.chat_history import BaraDBChatHistory
client = Client("localhost", 9472) client = Client("localhost", 9472)
@@ -59,7 +59,7 @@ class BaraDBChatHistory:
async def _ensure_table(self): async def _ensure_table(self):
if self._initialized: if self._initialized:
return return
await self.client.execute( await self.client.query(
f""" f"""
CREATE TABLE IF NOT EXISTS {self.table} ( CREATE TABLE IF NOT EXISTS {self.table} (
id TEXT PRIMARY KEY, id TEXT PRIMARY KEY,
@@ -73,7 +73,7 @@ class BaraDBChatHistory:
) )
""" """
) )
await self.client.execute( await self.client.query(
f"CREATE INDEX IF NOT EXISTS idx_{self.table}_session " f"CREATE INDEX IF NOT EXISTS idx_{self.table}_session "
f"ON {self.table}(session_id) USING btree" f"ON {self.table}(session_id) USING btree"
) )
@@ -98,14 +98,23 @@ class BaraDBChatHistory:
created_at = datetime.utcnow().isoformat() created_at = datetime.utcnow().isoformat()
for key, val in self._build_session().items(): for key, val in self._build_session().items():
await self.client.execute(f"SET {key} = '{val}'") await self.client.query_params(
f"SET {key} = $1", [self._wire_string(val)]
)
await self.client.execute( await self.client.query_params(
f"INSERT INTO {self.table} (id, session_id, role, content, metadata, " f"INSERT INTO {self.table} (id, session_id, role, content, metadata, tenant_id, user_id, created_at) "
f"tenant_id, user_id, created_at) " f"VALUES ($1, $2, $3, $4, $5, $6, $7, $8)",
f"VALUES ('{msg_id}', '{self.session_id}', '{role}', " [
f"'{_escape(content)}', '{_escape(metadata)}', " self._wire_string(msg_id),
f"'{self.tenant_id or ''}', '{self.user_id or ''}', '{created_at}')" self._wire_string(self.session_id),
self._wire_string(role),
self._wire_string(content),
self._wire_string(metadata),
self._wire_string(self.tenant_id or ""),
self._wire_string(self.user_id or ""),
self._wire_string(created_at),
],
) )
def add_user_message(self, message: Any) -> None: def add_user_message(self, message: Any) -> None:
@@ -132,14 +141,22 @@ class BaraDBChatHistory:
created_at = datetime.utcnow().isoformat() created_at = datetime.utcnow().isoformat()
for key, val in self._build_session().items(): for key, val in self._build_session().items():
await self.client.execute(f"SET {key} = '{val}'") await self.client.query_params(
f"SET {key} = $1", [self._wire_string(val)]
)
await self.client.execute( await self.client.query_params(
f"INSERT INTO {self.table} (id, session_id, role, content, " f"INSERT INTO {self.table} (id, session_id, role, content, tenant_id, user_id, created_at) "
f"tenant_id, user_id, created_at) " f"VALUES ($1, $2, $3, $4, $5, $6, $7)",
f"VALUES ('{msg_id}', '{self.session_id}', '{role}', " [
f"'{_escape(content)}', '{self.tenant_id or ''}', " self._wire_string(msg_id),
f"'{self.user_id or ''}', '{created_at}')" self._wire_string(self.session_id),
self._wire_string(role),
self._wire_string(content),
self._wire_string(self.tenant_id or ""),
self._wire_string(self.user_id or ""),
self._wire_string(created_at),
],
) )
async def get_messages(self) -> List[Any]: async def get_messages(self) -> List[Any]:
@@ -154,13 +171,19 @@ class BaraDBChatHistory:
return f"{self.type}: {self.content}" return f"{self.type}: {self.content}"
for key, val in self._build_session().items(): for key, val in self._build_session().items():
await self.client.execute(f"SET {key} = '{val}'") await self.client.query_params(
f"SET {key} = $1", [self._wire_string(val)]
)
result = await self.client.execute( result = await self.client.query_params(
f"SELECT role, content FROM {self.table} " f"SELECT role, content FROM {self.table} "
f"WHERE session_id = '{self.session_id}' " f"WHERE session_id = $1 "
f"ORDER BY created_at ASC " f"ORDER BY created_at ASC "
f"LIMIT {self.max_messages}" f"LIMIT $2",
[
self._wire_string(self.session_id),
self._wire_int(self.max_messages),
],
) )
messages = [] messages = []
if result and hasattr(result, "rows"): if result and hasattr(result, "rows"):
@@ -180,9 +203,12 @@ class BaraDBChatHistory:
async def clear(self) -> None: async def clear(self) -> None:
await self._ensure_table() await self._ensure_table()
for key, val in self._build_session().items(): for key, val in self._build_session().items():
await self.client.execute(f"SET {key} = '{val}'") await self.client.query_params(
await self.client.execute( f"SET {key} = $1", [self._wire_string(val)]
f"DELETE FROM {self.table} WHERE session_id = '{self.session_id}'" )
await self.client.query_params(
f"DELETE FROM {self.table} WHERE session_id = $1",
[self._wire_string(self.session_id)],
) )
async def get_session_summary(self, max_tokens: int = 2000) -> str: async def get_session_summary(self, max_tokens: int = 2000) -> str:
@@ -197,6 +223,13 @@ class BaraDBChatHistory:
total_chars += len(text) total_chars += len(text)
return "\n".join(parts) return "\n".join(parts)
@staticmethod
def _wire_string(val: str) -> Any:
# Lazy import to avoid circular dependency
from baradb import WireValue
return WireValue.string(val)
def _escape(s: str) -> str: @staticmethod
return s.replace("'", "''").replace("\\", "\\\\") def _wire_int(val: int) -> Any:
from baradb import WireValue
return WireValue.int64(val)
+68 -24
View File
@@ -2,7 +2,7 @@
BaraDB LangChain Vector Store Integration BaraDB LangChain Vector Store Integration
Usage: Usage:
from baradb import Client from baradb import Client, WireValue
from baradb.langchain_store import BaraDBStore from baradb.langchain_store import BaraDBStore
from langchain.embeddings import OpenAIEmbeddings from langchain.embeddings import OpenAIEmbeddings
@@ -50,6 +50,19 @@ class BaraDBStore:
self.vector_dimension = vector_dimension self.vector_dimension = vector_dimension
self._table_created = False self._table_created = False
def _wire(self, val: Any) -> Any:
"""Lazy import WireValue to avoid circular deps."""
from baradb import WireValue
if isinstance(val, str):
return WireValue.string(val)
if isinstance(val, int):
return WireValue.int64(val)
if isinstance(val, float):
return WireValue.float64(val)
if val is None:
return WireValue.null()
return WireValue.string(str(val))
async def _ensure_table(self) -> None: async def _ensure_table(self) -> None:
if self._table_created: if self._table_created:
return return
@@ -84,24 +97,23 @@ class BaraDBStore:
vec_str = "[" + ",".join(str(v) for v in vec) + "]" vec_str = "[" + ",".join(str(v) for v in vec) + "]"
meta = metadatas[i] if metadatas and i < len(metadatas) else {} meta = metadatas[i] if metadatas and i < len(metadatas) else {}
meta_cols = [] col_names = [self.embedding_col, self.text_col]
meta_vals = [] params = [self._wire(vec_str), self._wire(text)]
if self.tenant_id: if self.tenant_id:
meta_cols.append("tenant_id") col_names.append("tenant_id")
meta_vals.append(f"'{self.tenant_id}'") params.append(self._wire(self.tenant_id))
for mc in self.metadata_cols: for mc in self.metadata_cols:
if mc in meta: if mc in meta:
meta_cols.append(mc) col_names.append(mc)
meta_vals.append(f"'{meta[mc]}'") params.append(self._wire(meta[mc]))
col_list = f"{self.embedding_col}, {self.text_col}" placeholders = [f"${j + 1}" for j in range(len(params))]
val_list = f"'{vec_str}', '{text.replace(\"'\", \"''\")}'" sql = (
if meta_cols: f"INSERT INTO {self.table} ({', '.join(col_names)}) "
col_list += ", " + ", ".join(meta_cols) f"VALUES ({', '.join(placeholders)}) RETURNING id"
val_list += ", " + ", ".join(meta_vals) )
result = await self.client.query_params(sql, params)
sql = f"INSERT INTO {self.table} ({col_list}) VALUES ({val_list}) RETURNING id"
result = await self.client.query(sql)
if result.rows: if result.rows:
inserted_ids.append(result.rows[0].get("id", str(i))) inserted_ids.append(result.rows[0].get("id", str(i)))
else: else:
@@ -120,14 +132,36 @@ class BaraDBStore:
# Set tenant session variable if multi-tenant # Set tenant session variable if multi-tenant
if self.tenant_id: if self.tenant_id:
await self.client.query(f"SET app.tenant_id = '{self.tenant_id}'") await self.client.query_params(
"SET app.tenant_id = $1", [self._wire(self.tenant_id)]
)
if filter_col and filter_val: if filter_col and filter_val:
sql = f"SELECT hybrid_search_filtered('{self.table}', '{self.embedding_col}', '{self.text_col}', '{query.replace(\"'\", \"''\")}', '{vec_str}', {k}, '{filter_col}', '{filter_val}') AS res" sql = (
"SELECT hybrid_search_filtered($1, $2, $3, $4, $5, $6, $7, $8) AS res"
)
params = [
self._wire(self.table),
self._wire(self.embedding_col),
self._wire(self.text_col),
self._wire(query),
self._wire(vec_str),
self._wire(k),
self._wire(filter_col),
self._wire(filter_val),
]
else: else:
sql = f"SELECT hybrid_search('{self.table}', '{self.embedding_col}', '{self.text_col}', '{query.replace(\"'\", \"''\")}', '{vec_str}', {k}) AS res" sql = "SELECT hybrid_search($1, $2, $3, $4, $5, $6) AS res"
params = [
self._wire(self.table),
self._wire(self.embedding_col),
self._wire(self.text_col),
self._wire(query),
self._wire(vec_str),
self._wire(k),
]
result = await self.client.query(sql) result = await self.client.query_params(sql, params)
if not result.rows: if not result.rows:
return [] return []
@@ -141,8 +175,11 @@ class BaraDBStore:
for item in arr: for item in arr:
doc_id = item.get("id", "") doc_id = item.get("id", "")
score = float(item.get("score", 0)) score = float(item.get("score", 0))
# Fetch full row # Fetch full row — use parameterized query
row_result = await self.client.query(f"SELECT * FROM {self.table} WHERE id = {doc_id}") row_result = await self.client.query_params(
f"SELECT * FROM {self.table} WHERE id = $1",
[self._wire(doc_id)],
)
if row_result.rows: if row_result.rows:
page_content = row_result.rows[0].get(self.text_col, "") page_content = row_result.rows[0].get(self.text_col, "")
metadata = dict(row_result.rows[0]) metadata = dict(row_result.rows[0])
@@ -184,12 +221,19 @@ class BaraDBStore:
async def delete(self, ids: Optional[List[str]] = None) -> None: async def delete(self, ids: Optional[List[str]] = None) -> None:
await self._ensure_table() await self._ensure_table()
if ids: if ids:
id_list = ", ".join(str(i) for i in ids) # Build parameterized IN clause: $1, $2, ...
await self.client.query(f"DELETE FROM {self.table} WHERE id IN ({id_list})") placeholders = [f"${j + 1}" for j in range(len(ids))]
params = [self._wire(i) for i in ids]
await self.client.query_params(
f"DELETE FROM {self.table} WHERE id IN ({', '.join(placeholders)})",
params,
)
async def set_tenant(self, tenant_id: str) -> None: async def set_tenant(self, tenant_id: str) -> None:
self.tenant_id = tenant_id self.tenant_id = tenant_id
await self.client.query(f"SET app.tenant_id = '{tenant_id}'") await self.client.query_params(
"SET app.tenant_id = $1", [self._wire(tenant_id)]
)
class _SimpleDocument: class _SimpleDocument:
+22 -11
View File
@@ -181,6 +181,15 @@ class BaraDBClient:
resp = requests.post(f"{self.base}/query", json={"query": sql}, timeout=30) resp = requests.post(f"{self.base}/query", json={"query": sql}, timeout=30)
return resp.json() return resp.json()
def query_params(self, sql: str, params: list) -> dict:
"""Execute a parameterized query via HTTP API."""
resp = requests.post(
f"{self.base}/query",
json={"query": sql, "params": params},
timeout=30,
)
return resp.json()
def setup_bara_db(client: BaraDBClient, table: str = "rag_docs"): def setup_bara_db(client: BaraDBClient, table: str = "rag_docs"):
client.execute(f""" client.execute(f"""
@@ -216,16 +225,16 @@ def ingest_document(
chunk_idx = batch_start + i chunk_idx = batch_start + i
if embedding: if embedding:
vec_str = "[" + ",".join(str(v) for v in embedding) + "]" vec_str = "[" + ",".join(str(v) for v in embedding) + "]"
content_escaped = chunk.replace("'", "''") client.query_params(
client.execute(
f"INSERT INTO {table} (chunk_index, content, embedding) " f"INSERT INTO {table} (chunk_index, content, embedding) "
f"VALUES ({chunk_idx}, '{content_escaped}', '{vec_str}')" f"VALUES (? , ?, ?)",
[chunk_idx, chunk, vec_str],
) )
else: else:
content_escaped = chunk.replace("'", "''") client.query_params(
client.execute(
f"INSERT INTO {table} (chunk_index, content) " f"INSERT INTO {table} (chunk_index, content) "
f"VALUES ({chunk_idx}, '{content_escaped}')" f"VALUES (?, ?)",
[chunk_idx, chunk],
) )
print(f" Ingested chunks {batch_start + 1}-{min(batch_start + batch_size, len(chunks))}") print(f" Ingested chunks {batch_start + 1}-{min(batch_start + batch_size, len(chunks))}")
@@ -241,15 +250,17 @@ def search(
query_embedding = embed([query], embedder_config)[0] query_embedding = embed([query], embedder_config)[0]
if query_embedding: if query_embedding:
vec_str = "[" + ",".join(str(v) for v in query_embedding) + "]" vec_str = "[" + ",".join(str(v) for v in query_embedding) + "]"
result = client.execute( result = client.query_params(
f"SELECT id, chunk_index, content, cos_distance(embedding, '{vec_str}') AS distance " f"SELECT id, chunk_index, content, cos_distance(embedding, ?) AS distance "
f"FROM {table} " f"FROM {table} "
f"ORDER BY distance ASC " f"ORDER BY distance ASC "
f"LIMIT {k}" f"LIMIT ?",
[vec_str, k],
) )
else: else:
result = client.execute( result = client.query_params(
f"SELECT id, chunk_index, content FROM {table} LIMIT {k}" f"SELECT id, chunk_index, content FROM {table} LIMIT ?",
[k],
) )
if "rows" in result: if "rows" in result:
+2 -2
View File
@@ -55,7 +55,7 @@ proc embed*(e: Embedder, text: string): seq[float32] =
elif data.hasKey("data") and data["data"].kind == JArray and data["data"].len > 0: elif data.hasKey("data") and data["data"].kind == JArray and data["data"].len > 0:
for val in data["data"][0]["embedding"]: for val in data["data"][0]["embedding"]:
result.add(float32(val.getFloat())) result.add(float32(val.getFloat()))
except: except CatchableError:
discard discard
finally: finally:
client.close() client.close()
@@ -83,5 +83,5 @@ proc jsonToVector*(s: string): seq[float32] =
if p.len > 0: if p.len > 0:
try: try:
result.add(parseFloat(p)) result.add(parseFloat(p))
except: except CatchableError:
discard discard
+1 -1
View File
@@ -78,7 +78,7 @@ proc generate*(client: LLMClient, prompt: string, systemPrompt: string = ""): st
result = data["response"].getStr() result = data["response"].getStr()
elif data.hasKey("choices") and data["choices"].kind == JArray and data["choices"].len > 0: elif data.hasKey("choices") and data["choices"].kind == JArray and data["choices"].len > 0:
result = data["choices"][0]["message"]["content"].getStr() result = data["choices"][0]["message"]["content"].getStr()
except: except CatchableError:
result = "" result = ""
finally: finally:
httpClient.close() httpClient.close()
+1 -1
View File
@@ -171,7 +171,7 @@ proc parseCypher*(query: string): CypherQuery =
if limPos >= 0: if limPos >= 0:
try: try:
result.limit = parseInt(query[limPos + 7..^1].strip()) result.limit = parseInt(query[limPos + 7..^1].strip())
except: except CatchableError:
result.limit = 0 result.limit = 0
proc executeCypher*(g: Graph, query: CypherQuery): CypherResult = proc executeCypher*(g: Graph, query: CypherQuery): CypherResult =
+18 -11
View File
@@ -225,7 +225,7 @@ proc handleQuery(params: JsonNode): JsonNode =
var tokens: seq[qlexer.Token] var tokens: seq[qlexer.Token]
try: try:
tokens = qlexer.tokenize(sql) tokens = qlexer.tokenize(sql)
except: except CatchableError:
serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant
serverCtx.execCtx.currentUser = prevUser serverCtx.execCtx.currentUser = prevUser
return %*{"error": "Failed to tokenize SQL: " & getCurrentExceptionMsg()} return %*{"error": "Failed to tokenize SQL: " & getCurrentExceptionMsg()}
@@ -233,7 +233,7 @@ proc handleQuery(params: JsonNode): JsonNode =
var astNode: Node var astNode: Node
try: try:
astNode = qparser.parse(tokens) astNode = qparser.parse(tokens)
except: except CatchableError:
serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant
serverCtx.execCtx.currentUser = prevUser serverCtx.execCtx.currentUser = prevUser
return %*{"error": "Failed to parse SQL: " & getCurrentExceptionMsg()} return %*{"error": "Failed to parse SQL: " & getCurrentExceptionMsg()}
@@ -246,7 +246,7 @@ proc handleQuery(params: JsonNode): JsonNode =
var res: ExecResult var res: ExecResult
try: try:
res = executeQuery(serverCtx.execCtx, astNode, wireParams) res = executeQuery(serverCtx.execCtx, astNode, wireParams)
except: except CatchableError:
serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant
serverCtx.execCtx.currentUser = prevUser serverCtx.execCtx.currentUser = prevUser
return %*{"error": "Query execution failed: " & getCurrentExceptionMsg()} return %*{"error": "Query execution failed: " & getCurrentExceptionMsg()}
@@ -368,6 +368,7 @@ proc handleVectorSearch(params: JsonNode): JsonNode =
item["metadata"] = metaObj item["metadata"] = metaObj
jsonResults.add(item) jsonResults.add(item)
var prevTenant = serverCtx.execCtx.sessionVars.getOrDefault("app.tenant_id", "")
var sessionInfo = newJObject() var sessionInfo = newJObject()
if "tenant_id" in params and params["tenant_id"].kind == JString: if "tenant_id" in params and params["tenant_id"].kind == JString:
let tid = params["tenant_id"].getStr() let tid = params["tenant_id"].getStr()
@@ -375,7 +376,7 @@ proc handleVectorSearch(params: JsonNode): JsonNode =
sessionInfo["tenant_id"] = %tid sessionInfo["tenant_id"] = %tid
sessionInfo["user_id"] = %serverCtx.execCtx.currentUser sessionInfo["user_id"] = %serverCtx.execCtx.currentUser
return %*{ var r = %*{
"table": %table, "table": %table,
"column": %column, "column": %column,
"index_size": idx.nodes.len, "index_size": idx.nodes.len,
@@ -384,6 +385,8 @@ proc handleVectorSearch(params: JsonNode): JsonNode =
"results": jsonResults, "results": jsonResults,
"_session": sessionInfo "_session": sessionInfo
} }
serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant
return r
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Tool: schema_inspect # Tool: schema_inspect
@@ -394,6 +397,7 @@ proc handleSchemaInspect(params: JsonNode): JsonNode =
if params.kind == JObject and "table" in params and params["table"].kind == JString: if params.kind == JObject and "table" in params and params["table"].kind == JString:
targetTable = params["table"].getStr() targetTable = params["table"].getStr()
var prevTenant = serverCtx.execCtx.sessionVars.getOrDefault("app.tenant_id", "")
if "tenant_id" in params and params["tenant_id"].kind == JString: if "tenant_id" in params and params["tenant_id"].kind == JString:
serverCtx.execCtx.sessionVars["app.tenant_id"] = params["tenant_id"].getStr() serverCtx.execCtx.sessionVars["app.tenant_id"] = params["tenant_id"].getStr()
@@ -474,17 +478,20 @@ proc handleSchemaInspect(params: JsonNode): JsonNode =
jsonTables.add(tblObj) jsonTables.add(tblObj)
if targetTable.len > 0 and jsonTables.len == 0: if targetTable.len > 0 and jsonTables.len == 0:
serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant
return %*{"error": "Table '" & targetTable & "' not found"} return %*{"error": "Table '" & targetTable & "' not found"}
var sessionInfo = newJObject() var sessionInfo = newJObject()
sessionInfo["tenant_id"] = %serverCtx.execCtx.sessionVars.getOrDefault("app.tenant_id", "") sessionInfo["tenant_id"] = %serverCtx.execCtx.sessionVars.getOrDefault("app.tenant_id", "")
sessionInfo["user_id"] = %serverCtx.execCtx.currentUser sessionInfo["user_id"] = %serverCtx.execCtx.currentUser
return %*{ var r = %*{
"tables": jsonTables, "tables": jsonTables,
"table_count": jsonTables.len, "table_count": jsonTables.len,
"_session": sessionInfo "_session": sessionInfo
} }
serverCtx.execCtx.sessionVars["app.tenant_id"] = prevTenant
return r
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# MCP Protocol handlers # MCP Protocol handlers
@@ -578,14 +585,14 @@ proc writeToStdout(line: string) =
try: try:
stdout.writeLine(line) stdout.writeLine(line)
stdout.flushFile() stdout.flushFile()
except: except CatchableError:
discard discard
proc logToStderr*(msg: string) = proc logToStderr*(msg: string) =
try: try:
stderr.writeLine("[baradb-mcp] " & msg) stderr.writeLine("[baradb-mcp] " & msg)
stderr.flushFile() stderr.flushFile()
except: except CatchableError:
discard discard
proc processMessage(raw: string): string = proc processMessage(raw: string): string =
@@ -595,7 +602,7 @@ proc processMessage(raw: string): string =
var req: JsonNode var req: JsonNode
try: try:
req = parseJson(raw) req = parseJson(raw)
except: except CatchableError:
logToStderr("JSON parse error: " & getCurrentExceptionMsg()) logToStderr("JSON parse error: " & getCurrentExceptionMsg())
let resp = %*{ let resp = %*{
"jsonrpc": "2.0", "jsonrpc": "2.0",
@@ -640,7 +647,7 @@ proc processMessage(raw: string): string =
var dispResult: JsonNode var dispResult: JsonNode
try: try:
dispResult = dispatch(meth, params) dispResult = dispatch(meth, params)
except: except CatchableError:
logToStderr("Dispatch error for " & meth & ": " & getCurrentExceptionMsg()) logToStderr("Dispatch error for " & meth & ": " & getCurrentExceptionMsg())
let msg = getCurrentExceptionMsg() let msg = getCurrentExceptionMsg()
var errResp = %*{ var errResp = %*{
@@ -705,7 +712,7 @@ proc run*() =
if startupDone: if startupDone:
logToStderr("STDIN closed, exiting.") logToStderr("STDIN closed, exiting.")
break break
except: except CatchableError:
logToStderr("STDIN read error: " & getCurrentExceptionMsg()) logToStderr("STDIN read error: " & getCurrentExceptionMsg())
break break
@@ -740,7 +747,7 @@ when isMainModule:
try: try:
discard init(dataDir) discard init(dataDir)
run() run()
except: except CatchableError:
logToStderr("Fatal error: " & getCurrentExceptionMsg()) logToStderr("Fatal error: " & getCurrentExceptionMsg())
finally: finally:
close() close()
+9 -9
View File
@@ -1214,7 +1214,7 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
let textCol = evalExpr(expr.irFuncArgs[2], row, ctx) let textCol = evalExpr(expr.irFuncArgs[2], row, ctx)
let queryText = evalExpr(expr.irFuncArgs[3], row, ctx) let queryText = evalExpr(expr.irFuncArgs[3], row, ctx)
let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx) let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx)
let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except: 10 let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except ValueError: 10
let results = doHybridSearch(ctx, table, vecCol, textCol, queryText, queryVec, k) let results = doHybridSearch(ctx, table, vecCol, textCol, queryText, queryVec, k)
var parts: seq[string] = @[] var parts: seq[string] = @[]
for (id, score) in results: for (id, score) in results:
@@ -1227,7 +1227,7 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
let textCol = evalExpr(expr.irFuncArgs[2], row, ctx) let textCol = evalExpr(expr.irFuncArgs[2], row, ctx)
let queryText = evalExpr(expr.irFuncArgs[3], row, ctx) let queryText = evalExpr(expr.irFuncArgs[3], row, ctx)
let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx) let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx)
let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except: 10 let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except ValueError: 10
let results = doHybridSearch(ctx, table, vecCol, textCol, queryText, queryVec, k) let results = doHybridSearch(ctx, table, vecCol, textCol, queryText, queryVec, k)
var ids: seq[string] = @[] var ids: seq[string] = @[]
for (id, score) in results: for (id, score) in results:
@@ -1240,7 +1240,7 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
let textCol = evalExpr(expr.irFuncArgs[2], row, ctx) let textCol = evalExpr(expr.irFuncArgs[2], row, ctx)
let queryText = evalExpr(expr.irFuncArgs[3], row, ctx) let queryText = evalExpr(expr.irFuncArgs[3], row, ctx)
let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx) let queryVec = evalExpr(expr.irFuncArgs[4], row, ctx)
let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except: 10 let k = try: parseInt(evalExpr(expr.irFuncArgs[5], row, ctx)) except ValueError: 10
let filterCol = evalExpr(expr.irFuncArgs[6], row, ctx) let filterCol = evalExpr(expr.irFuncArgs[6], row, ctx)
let filterVal = evalExpr(expr.irFuncArgs[7], row, ctx) let filterVal = evalExpr(expr.irFuncArgs[7], row, ctx)
let results = doHybridSearchFiltered(ctx, table, vecCol, textCol, queryText, queryVec, k, filterCol, filterVal) let results = doHybridSearchFiltered(ctx, table, vecCol, textCol, queryText, queryVec, k, filterCol, filterVal)
@@ -1260,7 +1260,7 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
let queryTerms = queryText.toLower().splitWhitespace() let queryTerms = queryText.toLower().splitWhitespace()
for elem in arr: for elem in arr:
var score = 0.0 var score = 0.0
try: score = parseFloat(elem["score"].getStr()) except: discard try: score = parseFloat(elem["score"].getStr()) except ValueError: discard
# Simple term overlap boost # Simple term overlap boost
for term in queryTerms: for term in queryTerms:
if term.len > 2: if term.len > 2:
@@ -1274,16 +1274,16 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
for (elem, _) in boosted: for (elem, _) in boosted:
outArr.add(elem) outArr.add(elem)
return $(%* outArr) return $(%* outArr)
except: except CatchableError:
return resultsJson return resultsJson
of "chunk": of "chunk":
if expr.irFuncArgs.len < 1: return "[]" if expr.irFuncArgs.len < 1: return "[]"
let text = evalExpr(expr.irFuncArgs[0], row, ctx) let text = evalExpr(expr.irFuncArgs[0], row, ctx)
let maxSize = if expr.irFuncArgs.len >= 2: let maxSize = if expr.irFuncArgs.len >= 2:
try: parseInt(evalExpr(expr.irFuncArgs[1], row, ctx)) except: 1024 try: parseInt(evalExpr(expr.irFuncArgs[1], row, ctx)) except ValueError: 1024
else: 1024 else: 1024
let overlap = if expr.irFuncArgs.len >= 3: let overlap = if expr.irFuncArgs.len >= 3:
try: parseInt(evalExpr(expr.irFuncArgs[2], row, ctx)) except: 128 try: parseInt(evalExpr(expr.irFuncArgs[2], row, ctx)) except ValueError: 128
else: 128 else: 128
let cfg = chunkmod.ChunkConfig(maxChunkSize: maxSize, chunkOverlap: overlap, let cfg = chunkmod.ChunkConfig(maxChunkSize: maxSize, chunkOverlap: overlap,
strategy: chunkmod.csRecursive, minChunkSize: 64) strategy: chunkmod.csRecursive, minChunkSize: 64)
@@ -1426,7 +1426,7 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
if expr.irFuncArgs.len < 1: return "[]" if expr.irFuncArgs.len < 1: return "[]"
let graphName = evalExpr(expr.irFuncArgs[0], row, ctx) let graphName = evalExpr(expr.irFuncArgs[0], row, ctx)
let dims = if expr.irFuncArgs.len >= 2: let dims = if expr.irFuncArgs.len >= 2:
try: parseInt(evalExpr(expr.irFuncArgs[1], row, ctx)) except: 64 try: parseInt(evalExpr(expr.irFuncArgs[1], row, ctx)) except ValueError: 64
else: 64 else: 64
if graphName notin ctx.graphs: if graphName notin ctx.graphs:
return "[]" return "[]"
@@ -1514,7 +1514,7 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
var nodeId: int64 = 0 var nodeId: int64 = 0
if expr.irFuncArgs.len > 0: if expr.irFuncArgs.len > 0:
try: nodeId = parseInt(evalExpr(expr.irFuncArgs[0], row, ctx)) try: nodeId = parseInt(evalExpr(expr.irFuncArgs[0], row, ctx))
except: nodeId = 0 except ValueError: nodeId = 0
nodeId = nodeId and 0x3FF # 10 bits nodeId = nodeId and 0x3FF # 10 bits
let ts = int64(epochTime() * 1000) and 0x1FFFFFFFFFF # 41 bits let ts = int64(epochTime() * 1000) and 0x1FFFFFFFFFF # 41 bits
var snowSeq = int64(getMonoTime().ticks() and 0xFFF) # 12 bits from monotonic var snowSeq = int64(getMonoTime().ticks() and 0xFFF) # 12 bits from monotonic