Files
ragflow/rag/utils/gaussdb_conn.py
2026-08-18 09:53:11 +08:00

1432 lines
57 KiB
Python

#
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from __future__ import annotations
import json
import logging
import re
from math import isfinite
from typing import Any, Iterable
from psycopg2 import Error as PsycopgError, errorcodes
from pydantic import BaseModel
from common.constants import PAGERANK_FLD
from common.doc_store.doc_store_base import FusionExpr, MatchDenseExpr, MatchTextExpr
from common.doc_store.gaussdb_conn_base import (
GaussDBConnectionBase,
GaussDBSQLValidator,
GaussDBSearchBuilder,
extra_field_expr,
validate_extra_field,
)
logger = logging.getLogger("ragflow.gaussdb_conn")
SQL_QUERY_TIMEOUT_MS = 30000
VECTOR_COLUMN_RE = re.compile(r"^q_(?P<dim>\d+)_vec$")
VECTOR_VALID_COLUMN_RE = re.compile(r"^q_(?P<dim>\d+)_vec_valid$")
def _tokenize_query_terms(query_text: Any) -> list[str]:
from rag.nlp import rag_tokenizer
return [term for term in rag_tokenizer.tokenize(str(query_text or "")).split() if term]
def _application_highlight(text: Any, keywords: list[str]) -> str | None:
if text is None:
return None
simple_terms, ngram_terms = GaussDBSearchBuilder.split_text_query_terms(keywords)
unique_terms: dict[str, str] = {}
for term in [*simple_terms, *ngram_terms]:
unique_terms.setdefault(term.casefold(), term)
terms = sorted(unique_terms.values(), key=len, reverse=True)
if not terms:
return None
pattern = re.compile("|".join(re.escape(term) for term in terms), re.IGNORECASE)
highlighted, count = pattern.subn(lambda match: f"<em>{match.group(0)}</em>", str(text))
return highlighted if count else None
class GaussDBError(Exception):
pass
CHUNK_COLUMNS = (
"id",
"kb_id",
"doc_id",
"docnm_kwd",
"doc_type_kwd",
"title_tks",
"title_sm_tks",
"content_with_weight",
"content_ltks",
"content_sm_ltks",
"important_kwd",
"important_tks",
"question_kwd",
"question_tks",
"tag_kwd",
"tag_feas",
"available_int",
"pagerank_fea",
"create_time",
"create_timestamp_flt",
"img_id",
"position_int",
"page_num_int",
"top_int",
"metadata",
"chunk_data",
"extra",
"_order_id",
"group_id",
"mom_id",
"knowledge_graph_kwd",
"source_id",
"entity_kwd",
"entity_type_kwd",
"from_entity_kwd",
"to_entity_kwd",
"weight_int",
"weight_flt",
"entities_kwd",
"rank_flt",
"n_hop_with_weight",
"removed_kwd",
"raptor_kwd",
"raptor_layer_int",
)
CHUNK_COLUMN_SET = set(CHUNK_COLUMNS)
JSONB_COLUMNS = {
"important_kwd",
"question_kwd",
"tag_kwd",
"tag_feas",
"position_int",
"page_num_int",
"top_int",
"metadata",
"chunk_data",
"extra",
"source_id",
"entities_kwd",
}
JSONB_MULTI_VALUE_COLUMNS = {
"important_kwd",
"question_kwd",
"tag_kwd",
"source_id",
"entities_kwd",
}
KEY_COLUMNS = {"id", "kb_id"}
DOC_META_COLUMNS = ("id", "kb_id", "meta_fields")
DOC_META_COLUMN_SET = set(DOC_META_COLUMNS)
DOC_META_JSONB_COLUMNS = {"meta_fields"}
DEFAULT_VALUES = {
"available_int": 1,
"removed_kwd": "N",
}
_UNDEFINED_TABLE_SQLSTATE = "42P01"
def _is_undefined_table_error(exc: Exception) -> bool:
return getattr(exc, "pgcode", None) == _UNDEFINED_TABLE_SQLSTATE or getattr(exc, "sqlstate", None) == _UNDEFINED_TABLE_SQLSTATE
def _infer_chunk_vector_size(documents: list[dict]) -> int:
dimensions = {int(match.group("dim")) for document in documents for column in document if (match := VECTOR_COLUMN_RE.fullmatch(str(column)))}
if len(dimensions) != 1:
raise ValueError("cannot infer a single GaussDB vector dimension from chunk documents")
return dimensions.pop()
class SearchResult(BaseModel):
total: int
chunks: list[dict]
class GaussDBConnection(GaussDBConnectionBase):
def db_type(self) -> str:
return "gaussdb"
def create_idx(self, index_name: str, dataset_id: str, vector_size: int, parser_id: str = None):
# Keep the functional table even when the optional gsdiskann build
# fails, while still propagating the index error to the caller.
statements = [
self.ddl.build_advisory_lock_sql(f"create_idx:{self.schema}:{index_name}"),
self.ddl.build_chunk_table_ddl(index_name),
]
statements.extend(self.ddl.build_regular_index_ddls(index_name))
statements.append(self.ddl.build_fulltext_ugin_ddl(index_name))
statements.append(self.ddl.build_ngram_fulltext_ugin_ddl(index_name))
statements.extend(self.ddl.build_vector_column_ddls(index_name, vector_size))
self._execute_statements(statements)
self._create_diskann_index(index_name, vector_size)
return True
def _create_diskann_index(self, index_name: str, vector_size: int) -> None:
statements = [
self.ddl.build_advisory_lock_sql(f"create_vector_idx:{self.schema}:{index_name}:{vector_size}"),
"SET LOCAL maintenance_work_mem = '1GB'",
self.ddl.build_diskann_index_ddl(index_name, vector_size),
]
self._execute_statements(statements)
def create_doc_meta_idx(self, index_name: str):
if not is_doc_meta_table(index_name):
raise ValueError(f"invalid GaussDB document metadata table name: {index_name}")
statements = [
self.ddl.build_advisory_lock_sql(f"create_doc_meta_idx:{self.schema}:{index_name}"),
*self.ddl.build_doc_meta_table_ddls(index_name),
]
self._execute_statements(statements)
return True
def delete_idx(self, index_name: str, dataset_id: str | None):
table = self.ddl.qualified_name(index_name)
if dataset_id:
try:
self._execute_write(f"DELETE FROM {table} WHERE kb_id = %s", [dataset_id])
except PsycopgError as exc:
if exc.pgcode != errorcodes.UNDEFINED_TABLE:
raise
else:
self._execute_write(f"DROP TABLE IF EXISTS {table}", [])
def index_exist(self, index_name: str, dataset_id: str | None = None) -> bool:
self.ddl.validate_identifier(index_name)
row = self._fetch_one(
"""
SELECT 1
FROM information_schema.tables
WHERE table_schema = %s
AND table_name = %s
LIMIT 1
""",
[self.schema, index_name],
)
return bool(row)
def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None, refresh: str | bool = "wait_for") -> list[str]:
if not documents:
return []
document_ids = [str(document.get("id") or "") for document in documents]
try:
is_meta = is_doc_meta_table(index_name)
if is_meta:
sql, params = self._build_doc_meta_upsert(index_name, documents, knowledgebase_id)
else:
normalized, errors = self._normalize_chunk_rows(index_name, documents, knowledgebase_id)
if errors:
return [document_id for document_id in document_ids if document_id] or errors
sql, params = self._build_chunk_upsert(index_name, normalized)
try:
self._execute_write(sql, params, many=True)
except Exception as exc:
if is_meta or not _is_undefined_table_error(exc) or self.index_exist(index_name, knowledgebase_id):
raise
vector_size = _infer_chunk_vector_size(documents)
logger.info(
"GaussDB chunk table missing during insert; creating table=%s vector_size=%s",
index_name,
vector_size,
)
self.create_idx(index_name, knowledgebase_id, vector_size)
self._execute_write(sql, params, many=True)
return []
except Exception as exc:
logger.error("GaussDB insert failed for table=%s ids=%s error=%s", index_name, document_ids, exc)
return [document_id for document_id in document_ids if document_id] or [str(exc)]
def get(self, chunk_id: str, index_name: str, knowledgebase_ids: list[str]) -> dict | None:
if not chunk_id:
return None
is_meta = is_doc_meta_table(index_name)
kb_ids = normalize_kb_ids(knowledgebase_ids)
if not is_meta and not kb_ids:
return None
if not self.index_exist(index_name):
return None
table = self.ddl.qualified_name(index_name)
params = [chunk_id]
where = "id = %s"
if kb_ids:
placeholders = ", ".join(["%s"] * len(kb_ids))
where += f" AND kb_id IN ({placeholders})"
params.extend(kb_ids)
if not is_meta and len(kb_ids) > 1:
sql = f"SELECT * FROM {table} WHERE {where} ORDER BY kb_id ASC LIMIT 2"
rows, description = self._fetch_all_with_description(sql, params)
chunks = [self._row_to_chunk(row, description) for row in rows or []]
self._check_cross_kb_duplicate_chunk_ids(chunks)
return chunks[0] if chunks else None
sql = f"SELECT * FROM {table} WHERE {where} LIMIT 1"
row, description = self._fetch_one_with_description(sql, params)
if row is None:
return None
return self._row_to_chunk(row, description)
def search(
self,
select_fields: list[str],
highlight_fields: list[str],
condition: dict,
match_expressions: list,
order_by,
offset: int,
limit: int,
index_names: str | list[str],
knowledgebase_ids: list[str] | None = None,
agg_fields: list[str] | None = None,
rank_feature: dict | None = None,
**kwargs,
) -> SearchResult:
if knowledgebase_ids is None:
knowledgebase_ids = kwargs.get("dataset_ids") or []
tables = normalize_table_names(index_names)
if not tables:
return SearchResult(total=0, chunks=[])
if len(tables) == 1 and is_doc_meta_table(tables[0]) and not match_expressions and not agg_fields:
return self._search_doc_meta_table(
select_fields=select_fields,
condition=condition,
order_by=order_by,
offset=offset,
limit=limit,
index_name=tables[0],
knowledgebase_ids=knowledgebase_ids,
)
if any(is_doc_meta_table(table) for table in tables):
raise ValueError("document metadata tables cannot be mixed with chunk search")
if agg_fields:
return self._search_chunk_aggregation(
tables=tables,
condition=condition,
knowledgebase_ids=knowledgebase_ids,
agg_fields=agg_fields,
)
parsed = self._parse_match_expressions(match_expressions, rank_feature)
scoped_condition = self._scoped_search_condition(condition, knowledgebase_ids)
if len(tables) == 1:
return self._search_chunk_table(
table=tables[0],
select_fields=select_fields,
highlight_fields=highlight_fields,
condition=scoped_condition,
parsed=parsed,
order_by=order_by,
offset=offset,
limit=limit,
)
effective_offset = max(int(offset or 0), 0)
effective_limit = max(int(limit or 0), 0)
if effective_limit == 0:
effective_limit = 10000
collection_limit = effective_offset + effective_limit
total = 0
chunks: list[dict] = []
for table in tables:
result = self._search_chunk_table(
table=table,
select_fields=select_fields,
highlight_fields=highlight_fields,
condition=scoped_condition,
parsed=parsed,
order_by=order_by,
offset=0,
limit=collection_limit,
)
total += result.total
chunks.extend(result.chunks)
chunks = self._sort_search_chunks(chunks, order_by, has_match=bool(parsed["keywords"] or parsed["vector"]))
chunks = chunks[effective_offset : effective_offset + effective_limit]
return SearchResult(total=total, chunks=chunks)
def update(self, condition: dict, new_value: dict, index_name: str, knowledgebase_id: str) -> bool:
if not condition or not new_value:
return False
if not (normalize_kb_id(knowledgebase_id) or normalize_kb_id(condition.get("kb_id"))):
return False
try:
table = self.ddl.qualified_name(index_name)
is_meta = is_doc_meta_table(index_name)
dynamic_values = {validate_extra_field(key): value for key, value in new_value.items() if key not in {"add", "remove"} and is_dynamic_chunk_field(key, is_meta)}
if dynamic_values:
fixed_values = {key: value for key, value in new_value.items() if key not in dynamic_values}
return self._update_with_dynamic_values(
table=table,
condition=condition,
knowledgebase_id=knowledgebase_id,
fixed_values=fixed_values,
dynamic_values=dynamic_values,
)
set_sql, set_params = self._build_set_clause(new_value, is_meta=is_meta, condition=condition)
where_sql, where_params = self._build_where_clause(condition, knowledgebase_id, is_meta=is_meta)
sql = f"UPDATE {table} SET {set_sql} WHERE {where_sql}"
self._execute_write(sql, [*set_params, *where_params])
return True
except Exception as exc:
logger.error("GaussDB update failed for table=%s condition=%s error=%s", index_name, condition, exc)
return False
def _update_with_dynamic_values(
self,
*,
table: str,
condition: dict,
knowledgebase_id: str,
fixed_values: dict,
dynamic_values: dict[str, Any],
) -> bool:
supplied_extra = fixed_values.pop("extra", None)
set_sql, set_params = self._build_set_clause(
fixed_values,
is_meta=False,
condition=condition,
)
where_sql, where_params = self._build_where_clause(
condition,
knowledgebase_id,
is_meta=False,
)
if not where_sql:
return False
conn = self.pool.get_conn()
cur = None
try:
cur = conn.cursor()
cur.execute(
f"SELECT id, kb_id, extra FROM {table} WHERE {where_sql} FOR UPDATE",
where_params,
)
rows = cur.fetchall()
for row in rows or []:
row_id, kb_id, stored_extra = row
base_extra = supplied_extra if supplied_extra is not None else stored_extra
merged_extra = dict(parse_json_dict(base_extra))
merged_extra.update(dynamic_values)
assignments = [part for part in (set_sql, "extra = %s::jsonb") if part]
cur.execute(
f"UPDATE {table} SET {', '.join(assignments)} WHERE kb_id = %s AND id = %s",
[
*set_params,
json.dumps(merged_extra, ensure_ascii=False),
kb_id,
row_id,
],
)
conn.commit()
return True
except Exception:
conn.rollback()
raise
finally:
close_cursor(cur)
self.pool.put_conn(conn)
def delete(self, condition: dict, index_name: str, knowledgebase_id: str) -> int:
if not condition:
return 0
if not (normalize_kb_id(knowledgebase_id) or normalize_kb_id(condition.get("kb_id"))):
return 0
try:
table = self.ddl.qualified_name(index_name)
where_sql, where_params = self._build_where_clause(
condition,
knowledgebase_id,
is_meta=is_doc_meta_table(index_name),
)
return self._execute_write(f"DELETE FROM {table} WHERE {where_sql}", where_params)
except Exception as exc:
logger.error("GaussDB delete failed for table=%s condition=%s error=%s", index_name, condition, exc)
return 0
def fetch_metadata_doc_ids(
self,
index_name: str,
kb_ids: list[str],
sql_filter: str,
filter_params: list[Any],
limit: int,
) -> list[str]:
if not is_doc_meta_table(index_name):
raise ValueError(f"invalid GaussDB document metadata table name: {index_name}")
scoped_kb_ids = normalize_kb_ids(kb_ids)
if not scoped_kb_ids or not sql_filter:
return []
table = self.ddl.qualified_name(index_name)
placeholders = ", ".join(["%s"] * len(scoped_kb_ids))
effective_limit = limit if limit and limit > 0 else 10000
sql = f"SELECT id FROM {table} WHERE kb_id IN ({placeholders}) AND ({sql_filter}) ORDER BY id LIMIT %s"
rows = self._fetch_all(sql, [*scoped_kb_ids, *(filter_params or []), effective_limit])
doc_ids = []
for row in rows or []:
if isinstance(row, dict):
value = row.get("id")
elif isinstance(row, (list, tuple)) and row:
value = row[0]
else:
value = None
if value is not None:
doc_ids.append(str(value))
return doc_ids
def sql(self, sql: str, fetch_size: int = 128, format: str = "json"):
self.logger.debug("GaussDBConnection.sql get sql: %s", sql)
fetch_size = int(fetch_size or 128)
validated = GaussDBSQLValidator.readonly_guard(
default_limit=fetch_size,
execution_schema=self.schema,
).validate_and_patch(sql)
rows, description = self._fetch_all_with_description(
validated.sql,
[],
statement_timeout_ms=SQL_QUERY_TIMEOUT_MS,
)
columns = [desc[0] for desc in description]
def coerce_value(value):
if isinstance(value, bytes):
return value.decode("utf-8", errors="ignore")
if isinstance(value, (dict, list)):
return json.dumps(value, ensure_ascii=False)
return value
rows_list = [[coerce_value(value) for value in list(row)] for row in rows or []]
result = {
"columns": [{"name": column, "type": "text"} for column in columns],
"rows": rows_list,
}
if format == "markdown":
header = "|" + "|".join(columns) + "|" if columns else ""
separator = "|" + "|".join(["---" for _ in columns]) + "|" if columns else ""
body = "\n".join(["|" + "|".join([str(value) for value in row]) + "|" for row in rows_list])
result["markdown"] = "\n".join([line for line in [header, separator, body] if line])
return result
def get_total(self, res) -> int:
return int(res.total)
def get_doc_ids(self, res) -> list[str]:
return [row["id"] for row in res.chunks if "id" in row]
@staticmethod
def _check_cross_kb_duplicate_chunk_ids(rows) -> None:
seen_kb_ids = {}
for row in rows or []:
chunk_id = row.get("id")
if chunk_id is None:
continue
kb_id = row.get("kb_id")
if kb_id is None:
continue
previous_kb_id = seen_kb_ids.get(chunk_id)
if previous_kb_id is not None and previous_kb_id != kb_id:
raise GaussDBError(f"cross-KB duplicate chunk id: {chunk_id}")
seen_kb_ids[chunk_id] = kb_id
def get_fields(self, res, fields: list[str]) -> dict[str, dict]:
self._check_cross_kb_duplicate_chunk_ids(res.chunks)
result = {}
for row in res.chunks:
chunk_id = row.get("id")
if chunk_id is None:
continue
result[chunk_id] = {field: row[field] for field in fields if row.get(field) is not None}
return result
def get_highlight(self, res, keywords: list[str], field_name: str):
self._check_cross_kb_duplicate_chunk_ids(res.chunks)
highlights = {}
for row in res.chunks:
chunk_id = row.get("id")
if row.get("_highlight_source") is not None:
value = _application_highlight(row["_highlight_source"], keywords)
else:
value = row.get("_highlight") or row.get("highlight")
if chunk_id and value:
highlights[chunk_id] = value
return highlights
def get_aggregation(self, res, field_name: str):
counts = {}
result = []
for row in res.chunks:
if "value" in row and "count" in row:
result.append((row["value"], row["count"]))
continue
value = row.get(field_name)
if isinstance(value, list):
for item in value:
if isinstance(item, str) and item.strip():
counts[item] = counts.get(item, 0) + 1
elif isinstance(value, str) and value.strip():
counts[value] = counts.get(value, 0) + 1
return result or list(counts.items())
def get_scores(self, res) -> dict[str, float]:
self._check_cross_kb_duplicate_chunk_ids(res.chunks)
return {row["id"]: float(row.get("_score") or 0.0) for row in res.chunks if row.get("id") is not None}
def get_vector_dimensions(self, index_name: str) -> list[int]:
self.ddl.validate_identifier(index_name)
sql = """
SELECT column_name
FROM information_schema.columns
WHERE table_schema = %s AND table_name = %s
"""
rows = self._fetch_all(sql, [self.schema, index_name])
dims = []
for row in rows:
col = row[0] if isinstance(row, (list, tuple)) else row.get("column_name")
match = VECTOR_COLUMN_RE.fullmatch(str(col or ""))
if match:
dims.append(int(match.group("dim")))
return sorted(set(dims))
def adjust_chunk_pagerank_fea(
self,
chunk_id: str,
index_name: str,
dataset_id: str,
delta: int,
min_w: int = 0,
max_w: int = 100,
**kwargs,
) -> bool:
if not chunk_id or not dataset_id:
return False
table = self.ddl.qualified_name(index_name)
sql = f"UPDATE {table} SET pagerank_fea = GREATEST(%s, LEAST(%s, COALESCE(pagerank_fea, 0) + %s)) WHERE kb_id = %s AND id = %s"
return self._execute_write(sql, [min_w, max_w, delta, dataset_id, chunk_id]) > 0
def _normalize_chunk_rows(self, index_name: str, documents: list[dict], knowledgebase_id: str | None):
errors = []
rows = []
batch_dims = sorted({int(VECTOR_COLUMN_RE.fullmatch(key).group("dim")) for document in documents for key in document if VECTOR_COLUMN_RE.fullmatch(key)})
for document in documents:
doc_id = str(document.get("id") or "")
try:
rows.append(self._normalize_chunk_row(index_name, document, knowledgebase_id, batch_dims))
except Exception as exc:
logger.error("GaussDB normalize chunk failed id=%s error=%s", doc_id, exc)
errors.append(doc_id or str(exc))
return rows, errors
def _normalize_chunk_row(self, index_name: str, document: dict, knowledgebase_id: str | None, batch_dims: list[int]) -> dict:
chunk_id = document.get("id")
if not chunk_id:
raise ValueError("chunk id is required")
row = {"id": str(chunk_id)}
kb_id = normalize_kb_id(document.get("kb_id")) or knowledgebase_id
if not kb_id:
raise ValueError("kb_id is required")
if knowledgebase_id and kb_id != knowledgebase_id:
raise ValueError(f"kb_id {kb_id} does not match dataset_id {knowledgebase_id}")
row["kb_id"] = kb_id
extra = {}
for key, value in document.items():
if key in {"id", "kb_id"} or VECTOR_COLUMN_RE.fullmatch(key) or VECTOR_VALID_COLUMN_RE.fullmatch(key):
continue
if key not in CHUNK_COLUMN_SET:
extra[key] = value
continue
row[key] = normalize_column_value(key, value)
if extra:
existing_extra = row.get("extra")
if isinstance(existing_extra, str):
try:
existing_extra = json.loads(existing_extra)
except json.JSONDecodeError:
existing_extra = {}
if not isinstance(existing_extra, dict):
existing_extra = {}
existing_extra.update(extra)
row["extra"] = json.dumps(existing_extra, ensure_ascii=False)
metadata = row.get("metadata")
metadata_dict = parse_json_dict(metadata)
if metadata_dict:
if metadata_dict.get("_group_id"):
row["group_id"] = metadata_dict["_group_id"]
elif row.get("doc_id"):
row.setdefault("group_id", row.get("doc_id"))
if metadata_dict.get("_title"):
row["docnm_kwd"] = metadata_dict["_title"]
elif row.get("doc_id"):
row.setdefault("group_id", row.get("doc_id"))
vector_columns = [key for key in document if VECTOR_COLUMN_RE.fullmatch(key)]
if vector_columns:
for vector_col in vector_columns:
dim = int(VECTOR_COLUMN_RE.fullmatch(vector_col).group("dim"))
row[vector_col] = vector_literal(document[vector_col], dim)
row[self.ddl.vector_valid_column_name(dim)] = True
else:
dims = batch_dims or self.get_vector_dimensions(index_name)
if len(dims) != 1:
raise ValueError("cannot infer GaussDB vector dimension")
dim = dims[0]
row[self.ddl.vector_column_name(dim)] = zero_vector_literal(dim)
row[self.ddl.vector_valid_column_name(dim)] = False
for column, default in DEFAULT_VALUES.items():
row.setdefault(column, default)
return row
def _build_chunk_upsert(self, index_name: str, rows: list[dict]) -> tuple[str, list[list[Any]]]:
if not rows:
raise ValueError("rows are required")
columns = ordered_columns(rows)
table = self.ddl.qualified_name(index_name)
placeholders = ", ".join(self._placeholder(column) for column in columns)
update_columns = [column for column in columns if column not in {"id", "kb_id"}]
update_clause = ", ".join(f"{column} = VALUES({column})" for column in update_columns)
sql = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {update_clause}"
params = [[row.get(column) for column in columns] for row in rows]
return sql, params
def _build_doc_meta_upsert(
self,
index_name: str,
documents: list[dict],
knowledgebase_id: str | None,
) -> tuple[str, list[list[Any]]]:
table = self.ddl.qualified_name(index_name)
params = []
for document in documents:
doc_id = document.get("id")
kb_id = normalize_kb_id(document.get("kb_id")) or knowledgebase_id
if not doc_id or not kb_id:
raise ValueError("doc metadata id and kb_id are required")
if knowledgebase_id and kb_id != knowledgebase_id:
raise ValueError(f"kb_id {kb_id} does not match dataset_id {knowledgebase_id}")
params.append([str(doc_id), kb_id, json.dumps(document.get("meta_fields") or {}, ensure_ascii=False)])
sql = f"INSERT INTO {table} (id, kb_id, meta_fields) VALUES (%s, %s, %s::jsonb) ON DUPLICATE KEY UPDATE meta_fields = VALUES(meta_fields)"
return sql, params
def _search_doc_meta_table(
self,
select_fields: list[str],
condition: dict,
order_by,
offset: int,
limit: int,
index_name: str,
knowledgebase_ids: list[str],
) -> SearchResult:
effective_condition = dict(condition or {})
kb_ids = normalize_kb_ids(knowledgebase_ids)
condition_kb_ids = normalize_kb_ids(effective_condition.get("kb_id"))
if kb_ids:
if condition_kb_ids and not set(condition_kb_ids).issubset(set(kb_ids)):
raise ValueError("condition kb_id must stay within knowledgebase_ids")
if not condition_kb_ids:
effective_condition["kb_id"] = kb_ids
if not normalize_kb_ids(effective_condition.get("kb_id")):
raise ValueError("GaussDB document metadata search requires a kb_id boundary")
table = self.ddl.qualified_name(index_name)
columns = select_doc_meta_columns(select_fields)
where_sql, where_params = self._build_where_clause(effective_condition, None, is_meta=True)
sql = f"SELECT {', '.join(columns)}, COUNT(*) OVER() AS __total FROM {table}"
if where_sql:
sql += f" WHERE {where_sql}"
order_sql = build_doc_meta_order_by(order_by)
if order_sql:
sql += f" ORDER BY {order_sql}"
effective_limit = limit if limit and limit > 0 else 10000
effective_offset = max(int(offset or 0), 0)
sql += " LIMIT %s OFFSET %s"
rows, description = self._fetch_all_with_description(sql, [*where_params, effective_limit, effective_offset])
total = 0
chunks = []
for row in rows or []:
chunk = self._row_to_chunk(row, description)
total = int(chunk.pop("__total", total or 0) or 0)
chunks.append(chunk)
if not chunks and effective_offset:
count_sql = f"SELECT COUNT(*) FROM {table}"
if where_sql:
count_sql += f" WHERE {where_sql}"
row = self._fetch_one(count_sql, where_params)
total = int(row[0]) if row else 0
return SearchResult(total=total, chunks=chunks)
def _search_chunk_aggregation(
self,
tables: list[str],
condition: dict,
knowledgebase_ids: list[str] | None,
agg_fields: list[str],
) -> SearchResult:
if len(agg_fields) != 1:
raise ValueError("GaussDB search supports one aggregation field per request")
counts: dict[Any, int] = {}
scoped_condition = self._scoped_search_condition(condition, knowledgebase_ids)
for table in tables:
sql, params = self._search_builder().build_aggregation_sql(
table=table,
field_name=agg_fields[0],
condition=scoped_condition,
)
rows, description = self._fetch_all_with_description(sql, params)
for row in rows or []:
chunk = self._row_to_chunk(row, description)
if chunk.get("count") is not None:
chunk["count"] = int(chunk["count"])
value = chunk.get("value")
if value is not None:
counts[value] = counts.get(value, 0) + int(chunk.get("count") or 0)
chunks = [{"value": value, "count": count} for value, count in sorted(counts.items(), key=lambda item: (-item[1], str(item[0])))][:1000]
return SearchResult(total=len(chunks), chunks=chunks)
def _search_chunk_table(
self,
table: str,
select_fields: list[str],
highlight_fields: list[str],
condition: dict,
parsed: dict[str, Any],
order_by,
offset: int,
limit: int,
) -> SearchResult:
sql, params = self._search_builder().build_search_sql(
table=table,
select_fields=select_fields,
condition=condition,
keywords=parsed["keywords"],
vector=parsed["vector"],
vector_dim=parsed["vector_dim"],
vector_weight=parsed["vector_weight"],
similarity_threshold=parsed["similarity_threshold"],
topn=parsed["topn"],
offset=offset,
limit=limit,
highlight_fields=highlight_fields,
order_by=order_by,
pagerank_weight=parsed["pagerank_weight"],
)
rows, description = self._fetch_all_with_description(sql, params)
result = self._rows_to_search_result(rows, description)
if result.chunks or not max(int(offset or 0), 0):
return result
count_sql, count_params = self._search_builder().build_search_sql(
table=table,
select_fields=["id"],
condition=condition,
keywords=parsed["keywords"],
vector=parsed["vector"],
vector_dim=parsed["vector_dim"],
vector_weight=parsed["vector_weight"],
similarity_threshold=parsed["similarity_threshold"],
topn=parsed["topn"],
offset=0,
limit=1,
highlight_fields=[],
order_by=order_by,
pagerank_weight=parsed["pagerank_weight"],
)
count_rows, count_description = self._fetch_all_with_description(count_sql, count_params)
count_result = self._rows_to_search_result(count_rows, count_description)
return SearchResult(total=count_result.total, chunks=[])
def _rows_to_search_result(self, rows, description) -> SearchResult:
total = 0
chunks = []
for row in rows or []:
chunk = self._row_to_chunk(row, description)
total = int(chunk.pop("__total", total or 0) or 0)
chunks.append(chunk)
return SearchResult(total=total, chunks=chunks)
def _sort_search_chunks(self, chunks: list[dict], order_by, has_match: bool) -> list[dict]:
fields = getattr(order_by, "fields", None) or []
if not fields:
if has_match:
return sorted(
chunks,
key=lambda row: (
-(float(row.get("_score") or 0.0)),
str(row.get("kb_id") or ""),
str(row.get("id") or ""),
),
)
return sorted(chunks, key=lambda row: (str(row.get("kb_id") or ""), str(row.get("id") or "")))
sorted_chunks = list(chunks)
for field, direction in reversed(fields):
sorted_chunks.sort(
key=lambda row, field=field: sortable_search_value(row.get(field), field),
reverse=bool(direction),
)
return sorted_chunks
def _parse_match_expressions(
self,
match_expressions: list,
rank_feature: dict | None,
) -> dict[str, Any]:
keywords: list[str] = []
vector = None
vector_dim = None
topn = None
vector_weight = None
similarity_threshold = None
pagerank_weight = 10.0
for expr in match_expressions or []:
if isinstance(expr, MatchTextExpr):
query_text = (expr.extra_options or {}).get("original_query") or expr.matching_text or ""
keywords = _tokenize_query_terms(query_text)
topn = expr.topn if topn is None else min(topn, expr.topn)
elif isinstance(expr, MatchDenseExpr):
if expr.embedding_data_type != "float":
raise ValueError(f"unsupported GaussDB vector data type: {expr.embedding_data_type}")
vector = list(expr.embedding_data)
vector_dim = parse_vector_dim(expr.vector_column_name) or len(vector)
topn = expr.topn if topn is None else min(topn, expr.topn)
similarity_threshold = float((expr.extra_options or {}).get("similarity", 0.0))
elif isinstance(expr, FusionExpr):
vector_weight = parse_fusion_vector_weight(expr)
topn = expr.topn if topn is None else min(topn, expr.topn)
if rank_feature and rank_feature.get(PAGERANK_FLD) is not None:
pagerank_weight = float(rank_feature[PAGERANK_FLD])
if vector is None:
vector_weight = 0.0
elif not keywords:
vector_weight = 1.0 if vector_weight is None else float(vector_weight)
else:
vector_weight = 0.5 if vector_weight is None else float(vector_weight)
return {
"keywords": keywords,
"vector": vector,
"vector_dim": vector_dim,
"vector_weight": vector_weight,
"similarity_threshold": similarity_threshold,
"topn": topn,
"pagerank_weight": pagerank_weight,
}
def _scoped_search_condition(self, condition: dict | None, knowledgebase_ids: list[str] | None) -> dict:
effective = dict(condition or {})
kb_ids = normalize_kb_ids(knowledgebase_ids)
if kb_ids:
condition_kb_ids = normalize_kb_ids(effective.get("kb_id"))
if condition_kb_ids:
if not set(condition_kb_ids).issubset(set(kb_ids)):
raise ValueError("condition kb_id must stay within knowledgebase_ids")
else:
effective["kb_id"] = kb_ids
if "doc_ids" in effective and "doc_id" not in effective:
effective["doc_id"] = effective.pop("doc_ids")
if isinstance(effective.get("doc_id"), (list, tuple, set)) and not effective["doc_id"]:
effective.pop("doc_id")
if not normalize_kb_ids(effective.get("kb_id")):
raise ValueError("GaussDB chunk search requires a kb_id boundary")
return effective
def _search_builder(self) -> GaussDBSearchBuilder:
builder = getattr(self, "search_builder", None)
if builder is None:
builder = GaussDBSearchBuilder(schema=self.schema)
self.search_builder = builder
return builder
def _build_set_clause(self, new_value: dict, is_meta: bool, condition: dict) -> tuple[str, list[Any]]:
fragments = []
params = []
allowed_columns = DOC_META_COLUMN_SET if is_meta else CHUNK_COLUMN_SET
jsonb_columns = DOC_META_JSONB_COLUMNS if is_meta else JSONB_COLUMNS
jsonb_operations: dict[str, dict[str, Any]] = {}
updated_vector_dims = {int(match.group("dim")) for key in new_value if not is_meta and (match := VECTOR_COLUMN_RE.fullmatch(str(key)))}
for operation, value in new_value.items():
if operation in {"remove", "add"} and isinstance(value, dict):
for column, item in value.items():
jsonb_operations.setdefault(column, {})[operation] = item
emitted_jsonb_columns = set()
def append_jsonb_update(column: str) -> None:
if column in emitted_jsonb_columns:
return
emitted_jsonb_columns.add(column)
operations = jsonb_operations[column]
remove_expr = None
remove_param = None
if "remove" in operations:
escaped_item = str(operations["remove"]).replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
remove_expr = f"json_remove({column}, json_unquote(json_search({column}, 'one', jsonb_array_element_text(%s::jsonb, 0))))::jsonb"
remove_param = json.dumps([escaped_item], ensure_ascii=False)
if "add" in operations:
base_expr = f"COALESCE({remove_expr}, {column}, '[]'::jsonb)" if remove_expr is not None else f"COALESCE({column}, '[]'::jsonb)"
expression = (
f"(SELECT CASE WHEN candidate @> %s::jsonb THEN candidate ELSE json_array_append(candidate, '$', %s::jsonb)::jsonb END FROM (SELECT {base_expr} AS candidate) AS jsonb_update)"
)
params.extend(
[
json.dumps([operations["add"]], ensure_ascii=False),
json.dumps(operations["add"], ensure_ascii=False),
]
)
if remove_param is not None:
params.append(remove_param)
else:
expression = f"COALESCE({remove_expr}, {column})"
params.append(remove_param)
fragments.append(f"{column} = {expression}")
for key, value in new_value.items():
if key == "remove":
if isinstance(value, str):
if value in KEY_COLUMNS:
raise ValueError(f"key column cannot be updated: {value}")
if value not in allowed_columns:
raise ValueError(f"unsupported remove target: {value}")
fragments.append(f"{value} = NULL")
continue
if not isinstance(value, dict):
raise ValueError(f"unsupported remove target: {value}")
for column, item in value.items():
if column not in JSONB_MULTI_VALUE_COLUMNS:
raise ValueError(f"unsupported JSONB remove target: {column}")
append_jsonb_update(column)
continue
if key == "add":
if not isinstance(value, dict):
raise ValueError(f"unsupported add target: {value}")
for column, item in value.items():
if column not in JSONB_MULTI_VALUE_COLUMNS:
raise ValueError(f"unsupported JSONB add target: {column}")
append_jsonb_update(column)
continue
if key == "metadata" and not is_meta and isinstance(value, dict):
fragments.append("metadata = %s::jsonb")
params.append(json.dumps(value, ensure_ascii=False))
if value.get("_group_id"):
fragments.append("group_id = %s")
params.append(value["_group_id"])
if value.get("_title"):
fragments.append("docnm_kwd = %s")
params.append(value["_title"])
continue
if key in KEY_COLUMNS:
if key == "id" and isinstance(value, str) and value == condition.get("id"):
continue
raise ValueError(f"key column cannot be updated: {key}")
if vector_match := VECTOR_COLUMN_RE.fullmatch(key):
dim = int(vector_match.group("dim"))
fragments.append(f"{key} = {self._placeholder(key)}")
params.append(vector_literal(value, dim))
fragments.append(f"{self.ddl.vector_valid_column_name(dim)} = TRUE")
continue
if valid_match := VECTOR_VALID_COLUMN_RE.fullmatch(str(key)):
dim = self.ddl.validate_vector_dim(int(valid_match.group("dim")))
if dim in updated_vector_dims:
continue
fragments.append(f"{key} = %s")
params.append(bool(value))
continue
if key not in allowed_columns:
raise ValueError(f"unknown column for update: {key}")
if key in jsonb_columns:
fragments.append(f"{key} = %s::jsonb")
params.append(json.dumps(value, ensure_ascii=False))
else:
fragments.append(f"{key} = %s")
params.append(value)
return ", ".join(fragments), params
def _build_where_clause(self, condition: dict, knowledgebase_id: str | None, is_meta: bool) -> tuple[str, list[Any]]:
effective = dict(condition or {})
if knowledgebase_id:
existing_kb = normalize_kb_id(effective.get("kb_id"))
if existing_kb and existing_kb != knowledgebase_id:
raise ValueError(f"condition kb_id {existing_kb} does not match dataset_id {knowledgebase_id}")
effective["kb_id"] = knowledgebase_id
fragments = []
params = []
def column_expression(column: str) -> str:
validate_filter_column(column, is_meta)
return filter_column_expr(column, is_meta)
for key, value in effective.items():
if key == "exists":
column = column_expression(value)
if is_jsonb_dynamic_chunk_field(value, is_meta):
fragments.append(f"({column} IS NOT NULL AND {column} <> 'null'::jsonb)")
else:
fragments.append(f"{column} IS NOT NULL")
continue
if key == "must_not" and isinstance(value, dict) and "exists" in value:
column = column_expression(value["exists"])
if is_jsonb_dynamic_chunk_field(value["exists"], is_meta):
fragments.append(f"({column} IS NULL OR {column} = 'null'::jsonb)")
else:
fragments.append(f"{column} IS NULL")
continue
column = column_expression(key)
if is_jsonb_dynamic_chunk_field(key, is_meta):
values = list(value) if isinstance(value, (list, tuple, set)) else [value]
if not values:
raise ValueError(f"empty list condition for {key}")
predicates = []
for item in values:
predicates.append(f"({column} = %s::jsonb OR {column} @> %s::jsonb)")
params.extend(
[
json.dumps(item, ensure_ascii=False),
json.dumps([item], ensure_ascii=False),
]
)
fragments.append("(" + " OR ".join(predicates) + ")")
continue
if key in JSONB_MULTI_VALUE_COLUMNS:
values = list(value) if isinstance(value, (list, tuple, set)) else [value]
if not values:
raise ValueError(f"empty list condition for {key}")
fragments.append("(" + " OR ".join([f"{column} @> %s::jsonb"] * len(values)) + ")")
params.extend(json.dumps([item], ensure_ascii=False) for item in values)
elif isinstance(value, (list, tuple, set)):
values = list(value)
if not values:
raise ValueError(f"empty list condition for {key}")
fragments.append(f"{column} IN ({', '.join(['%s'] * len(values))})")
params.extend(values)
else:
fragments.append(f"{column} = %s")
params.append(value)
return " AND ".join(fragments), params
def _placeholder(self, column: str) -> str:
match = VECTOR_COLUMN_RE.fullmatch(column)
if match:
return f"%s::floatvector({match.group('dim')})"
if column in JSONB_COLUMNS or column in DOC_META_JSONB_COLUMNS:
return "%s::jsonb"
return "%s"
def _row_to_chunk(self, row, description) -> dict:
columns = [desc[0] for desc in description]
raw = dict(zip(columns, row)) if not isinstance(row, dict) else dict(row)
result = {}
dynamic_fields = {}
invalid_vectors = {f"q_{match.group('dim')}_vec" for key, value in raw.items() if (match := VECTOR_VALID_COLUMN_RE.fullmatch(key)) and value is False}
for key, value in raw.items():
if key in invalid_vectors:
continue
decoded = decode_column_value(key, value)
result[key] = decoded
if key == "extra" and isinstance(decoded, dict):
dynamic_fields.update(decoded)
for key, value in dynamic_fields.items():
result.setdefault(key, value)
return result
def _execute_statements(self, statements: Iterable[str | tuple[str, list[Any]]]) -> None:
conn = self.pool.get_conn()
cur = None
try:
cur = conn.cursor()
for statement in statements:
if isinstance(statement, tuple):
cur.execute(statement[0], statement[1])
else:
cur.execute(statement)
conn.commit()
except Exception:
conn.rollback()
raise
finally:
close_cursor(cur)
self.pool.put_conn(conn)
def _execute_write(self, sql: str, params: list[Any], many: bool = False) -> int:
conn = self.pool.get_conn()
cur = None
try:
cur = conn.cursor()
if many:
cur.executemany(sql, params)
else:
cur.execute(sql, params)
conn.commit()
return int(getattr(cur, "rowcount", 0) or 0)
except Exception:
conn.rollback()
raise
finally:
close_cursor(cur)
self.pool.put_conn(conn)
def _fetch_one(self, sql: str, params: list[Any]):
row, _description = self._fetch_one_with_description(sql, params)
return row
def _fetch_one_with_description(self, sql: str, params: list[Any]):
conn = self.pool.get_conn()
cur = None
try:
cur = conn.cursor()
cur.execute(sql, params)
return cur.fetchone(), getattr(cur, "description", None) or []
finally:
close_cursor(cur)
self.pool.put_conn(conn)
def _fetch_all(self, sql: str, params: list[Any]):
conn = self.pool.get_conn()
cur = None
try:
cur = conn.cursor()
cur.execute(sql, params)
return cur.fetchall()
finally:
close_cursor(cur)
self.pool.put_conn(conn)
def _fetch_all_with_description(self, sql: str, params: list[Any], statement_timeout_ms: int | None = None):
conn = self.pool.get_conn()
cur = None
try:
cur = conn.cursor()
if statement_timeout_ms and statement_timeout_ms > 0:
cur.execute(f"SET LOCAL statement_timeout = {int(statement_timeout_ms)}")
cur.execute(sql, params)
return cur.fetchall(), getattr(cur, "description", None) or []
finally:
close_cursor(cur)
self.pool.put_conn(conn)
def is_doc_meta_table(index_name: str) -> bool:
return str(index_name or "").startswith("ragflow_doc_meta_")
def normalize_kb_id(value) -> str | None:
if isinstance(value, list):
return str(value[0]) if value else None
if value is None:
return None
return str(value)
def normalize_kb_ids(values) -> list[str]:
if values is None:
return []
if isinstance(values, (str, bytes)):
raw_values = [values]
else:
raw_values = list(values)
normalized = []
seen = set()
for value in raw_values:
if value in (None, ""):
continue
text = str(value)
if text in seen:
continue
seen.add(text)
normalized.append(text)
return normalized
def normalize_table_names(index_names: str | list[str]) -> list[str]:
if isinstance(index_names, str):
return [name.strip() for name in index_names.split(",") if name.strip()]
return [str(name).strip() for name in index_names or [] if str(name).strip()]
def parse_vector_dim(vector_column_name: str) -> int | None:
match = VECTOR_COLUMN_RE.fullmatch(str(vector_column_name or ""))
return int(match.group("dim")) if match else None
def parse_fusion_vector_weight(expr: FusionExpr) -> float | None:
params = expr.fusion_params or {}
weights = params.get("weights")
if not weights:
return None
try:
return float(str(weights).split(",")[-1])
except (TypeError, ValueError):
return None
def sortable_search_value(value: Any, field: str):
if field == "page_num_int":
return nested_numeric_value(value, [0])
if field == "position_int":
return nested_numeric_value(value, [0, 3])
if field == "top_int":
return nested_numeric_value(value, [0])
if isinstance(value, (int, float)):
return value
if value is None:
return ""
return str(value)
def nested_numeric_value(value: Any, path: list[int]) -> float:
current = value
try:
for index in path:
current = current[index]
numeric = float(current)
return numeric if isfinite(numeric) else 100000000.0
except (TypeError, ValueError, IndexError, KeyError, OverflowError):
return 100000000.0
def select_doc_meta_columns(select_fields: list[str]) -> list[str]:
if not select_fields or "*" in select_fields:
return list(DOC_META_COLUMNS)
columns = []
for field in select_fields:
validate_filter_column(field, is_meta=True)
if field not in columns:
columns.append(field)
return columns
def build_doc_meta_order_by(order_by) -> str:
fields = getattr(order_by, "fields", None) or []
parts = []
for field, direction in fields:
validate_filter_column(field, is_meta=True)
parts.append(f"{field} {'DESC' if direction else 'ASC'}")
return ", ".join(parts)
def normalize_column_value(column: str, value: Any) -> Any:
if value is None:
return None
if column in JSONB_COLUMNS:
return json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
return value
def parse_json_dict(value) -> dict:
if isinstance(value, dict):
return value
if isinstance(value, str):
try:
parsed = json.loads(value)
except json.JSONDecodeError:
return {}
return parsed if isinstance(parsed, dict) else {}
return {}
def vector_literal(value, dim: int) -> str:
if isinstance(value, str):
parsed = parse_vector_literal(value)
if parsed is None:
raise ValueError("invalid vector literal")
value = parsed
if not isinstance(value, (list, tuple)) or len(value) != dim:
raise ValueError(f"vector dimension mismatch: expected {dim}, got {len(value) if hasattr(value, '__len__') else 'unknown'}")
return "[" + ",".join(str(item) for item in value) + "]"
def parse_vector_literal(value: str):
try:
parsed = json.loads(value)
except json.JSONDecodeError:
return None
return parsed if isinstance(parsed, list) else None
def zero_vector_literal(dim: int) -> str:
return "[" + ",".join(["0"] * dim) + "]"
def ordered_columns(rows: list[dict]) -> list[str]:
present = set()
for row in rows:
present.update(row)
dynamic_vectors = sorted(column for column in present if VECTOR_COLUMN_RE.fullmatch(column))
dynamic_valid = sorted(column for column in present if VECTOR_VALID_COLUMN_RE.fullmatch(column))
ordered = [column for column in CHUNK_COLUMNS if column in present]
ordered.extend(column for column in dynamic_vectors + dynamic_valid if column not in ordered)
return ordered
def validate_filter_column(column: str, is_meta: bool) -> None:
if is_meta:
if column not in DOC_META_COLUMN_SET:
raise ValueError(f"unsupported metadata filter column: {column}")
return
if column not in CHUNK_COLUMN_SET and not VECTOR_COLUMN_RE.fullmatch(str(column)) and not VECTOR_VALID_COLUMN_RE.fullmatch(str(column)):
validate_extra_field(column)
def is_dynamic_chunk_field(column: str, is_meta: bool = False) -> bool:
return not is_meta and column not in CHUNK_COLUMN_SET and not VECTOR_COLUMN_RE.fullmatch(str(column)) and not VECTOR_VALID_COLUMN_RE.fullmatch(str(column))
def is_jsonb_dynamic_chunk_field(column: str, is_meta: bool = False) -> bool:
return is_dynamic_chunk_field(column, is_meta) and column not in GaussDBSearchBuilder.JSONB_EXTRA_SCALAR_COLUMNS
def filter_column_expr(column: str, is_meta: bool = False) -> str:
validate_filter_column(column, is_meta)
if not is_dynamic_chunk_field(column, is_meta):
return str(column)
if column in GaussDBSearchBuilder.JSONB_EXTRA_SCALAR_COLUMNS:
return f"(extra #>> '{{{column}}}')"
return extra_field_expr(column)
def decode_column_value(column: str, value: Any) -> Any:
if value is None:
return None
if column in JSONB_COLUMNS or column in DOC_META_JSONB_COLUMNS:
if isinstance(value, str):
try:
return json.loads(value)
except json.JSONDecodeError as exc:
raise ValueError(f"invalid JSONB value for {column}") from exc
return value
if VECTOR_COLUMN_RE.fullmatch(column):
if isinstance(value, str):
parsed = parse_vector_literal(value)
if parsed is None:
raise ValueError(f"invalid vector value for {column}")
return parsed
return value
return value
def close_cursor(cur) -> None:
if cur is not None and hasattr(cur, "close"):
cur.close()