Skip to content

feat: add OceanBase/seekdb vector and graph storage providers#2130

Open
Evenss wants to merge 2 commits into
MemTensor:mainfrom
Evenss:main
Open

feat: add OceanBase/seekdb vector and graph storage providers#2130
Evenss wants to merge 2 commits into
MemTensor:mainfrom
Evenss:main

Conversation

@Evenss

@Evenss Evenss commented Jul 20, 2026

Copy link
Copy Markdown

Description

Add two optional storage providers for OceanBase / seekdb, reusing the existing BaseVecDB / BaseGraphDB contracts. No default behavior changes: unless a user selects the oceanbase / seekdb backend, everything works exactly as before.

  • OceanBaseVecDB (vec_dbs/oceanbase.py): built on pyseekdb's Collection / vector API, serving General Memory. Config now enforces a positive vector_dimension.
  • OceanBaseGraphDB (graph_dbs/oceanbase.py): ported from the PostgreSQL + pgvector backend (nodes + edges + JSON properties + VECTOR column) over the MySQL-compatible protocol. Includes a thread-safe connection pool (maxconn now effective), atomic multi-step deletes, and identifier whitelisting for table_prefix / search_filter keys.
  • Register oceanbase / seekdb aliases in the vector & graph factories and config factories; add the GraphDBError exception.

Dependencies: adds an optional extra ob-mem (containing pyseekdb); imports are guarded with try/except ImportError and are not added to core dependencies.

Related Issue (Required): Fixes #2109

Type of change

  • New feature (non-breaking change which adds functionality)

How Has This Been Tested?

  • Unit Test
.venv/bin/python -m pytest tests/vec_dbs/test_oceanbase.py tests/graph_dbs/test_oceanbase.py -q
# 31 passed

.venv/bin/ruff check src/memos/vec_dbs/oceanbase.py src/memos/graph_dbs/oceanbase.py
.venv/bin/ruff format --check src/memos/vec_dbs/oceanbase.py src/memos/graph_dbs/oceanbase.py
# All checks passed

External seekdb / OceanBase drivers (pyseekdb / pymysql) are stubbed with mocks; no live server is required. A live-server smoke test is not included in this PR.

Checklist

  • I have performed a self-review of my own code | 我已自行检查了自己的代码
  • I have commented my code in hard-to-understand areas | 我已在难以理解的地方对代码进行了注释
  • I have added tests that prove my fix is effective or that my feature works | 我已添加测试以证明我的修复有效或功能正常
  • I have created related documentation issue/PR in MemOS-Docs (if applicable) | 我已在 MemOS-Docs 中创建了相关的文档 issue/PR(如果适用)
  • I have linked the issue to this PR (if applicable) | 我已将 issue 链接到此 PR(如果适用)
  • I have mentioned the person who will review this PR | 我已提及将审查此 PR 的人

Reviewer Checklist

Add optional OceanBase / seekdb backends reusing the existing BaseVecDB
and BaseGraphDB contracts, without changing any default behavior.

- vec_dbs/oceanbase.py: OceanBaseVecDB on top of pyseekdb's Collection API,
  serving General Memory; require a positive vector_dimension in config.
- graph_dbs/oceanbase.py: OceanBaseGraphDB ported from the postgres backend
  (nodes + edges + JSON + VECTOR) over the MySQL-compatible protocol, with a
  thread-safe connection pool, atomic multi-step deletes, and identifier
  whitelisting (table_prefix / search_filter keys).
- Register "oceanbase" / "seekdb" aliases in the vec/graph factories and
  config factories; add GraphDBError; declare the optional "ob-mem" extra.
- Add contract tests for both providers.
@Memtensor-AI Memtensor-AI added area:api 云服务 / FastAPI / OpenAPI / MCP area:core MOS 编排层 / 框架底座 / 跨模块问题 area:database graph_db + vector_db | 图数据库与向量数据库 status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Jul 20, 2026
@Memtensor-AI

Memtensor-AI commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

🤖 Open Code Review

Target: PR #2130
Task: 886e3a9bdb3a8bb9
Base: main
Head: main

🔍 OpenCodeReview found 12 issue(s) in this PR.

⚠️ 1 warning(s) occurred during review.


1. pyproject.toml (L117)

The pyseekdb version pin >=1.4.0,<1.5.0 (patch-level upper bound) is significantly more restrictive than every other optional dependency in this file, which use minor-level upper bounds (e.g. <2.0.0, <3.0.0, <6.0.0). If pyseekdb releases a 1.4.1 or 1.4.2 patch that is otherwise compatible, users will be blocked from adopting it without a MemOS release. Using >=1.4.0,<2.0.0 (or at least <1.5.0 only if a known breaking change is expected at 1.5) would be more permissive and consistent with the rest of the file.

💡 Suggested Change

Before:

    "pyseekdb (>=1.4.0,<1.5.0)",  # Unified client for seekdb / OceanBase (Collection + raw SQL)

After:

    "pyseekdb (>=1.4.0,<2.0.0)",  # Unified client for seekdb / OceanBase (Collection + raw SQL)

2. src/memos/configs/vec_db.py (L66-L72)

In Pydantic v2, @field_validator does not run on default values — it is only invoked when the field is explicitly provided in the input data. Because vector_dimension inherits default=None from BaseVecDBConfig, a user who simply omits vector_dimension from their OceanBase config will bypass this validator entirely. None will silently reach HNSWConfiguration(dimension=None, ...) at runtime and produce a confusing low-level error rather than a clear config validation failure.

To also validate the default, add always=True (Pydantic v2 equivalent via @field_validator(..., always=True) — or more robustly, move the check into a @model_validator(mode='after')):

@model_validator(mode="after")
def validate_vector_dimension(self) -> "OceanBaseVecDBConfig":
    """OceanBase requires a concrete vector dimension to build the HNSW index."""
    value = self.vector_dimension
    if value is None or isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError("`vector_dimension` must be a positive integer for OceanBase")
    return self

3. src/memos/configs/graph_db.py (L311-L312)

When use_multi_db=True, user_name is left unvalidated. But in OceanBaseGraphDB, virtually every data operation resolves the active tenant via user_name = user_name or self.user_name, where self.user_name = config.user_name. If user_name is None in multi-db mode and no per-call user_name is passed, all tenant-scoped SQL filters (WHERE user_name = %s) are silently skipped, causing data from all tenants to be returned or modified together.

Either require user_name unconditionally, or document clearly that callers are responsible for always supplying it when use_multi_db=True.

💡 Suggested Change

Before:

        if not self.use_multi_db and not self.user_name:
            raise ValueError("In single-database mode, `user_name` must be provided")

After:

        if not self.user_name:
            raise ValueError("`user_name` must be provided for data isolation")

4. src/memos/vec_dbs/oceanbase.py (L148-L154)

The pyseekdb (chromadb-style) query() API expects query_embeddings to be a list of vectors (a batch), not a single flat vector. The response parsing already handles the batched shape correctly — it takes [0] of each result list — but the input is passed as a flat list rather than wrapped in one. This will likely cause a runtime error or silently search against the wrong data depending on how pyseekdb validates its input.

Change query_embeddings=query_vector to query_embeddings=[query_vector].

💡 Suggested Change

Before:

        response = self.collection.query(
            query_embeddings=query_vector,
            n_results=top_k,
            where=filter or None,
            include=["metadatas", "embeddings", "distances"],
        )
        ids = (response.get("ids") or [[]])[0]

After:

        response = self.collection.query(
            query_embeddings=[query_vector],
            n_results=top_k,
            where=filter or None,
            include=["metadatas", "embeddings", "distances"],
        )
        ids = (response.get("ids") or [[]])[0]

5. tests/graph_dbs/test_oceanbase.py (L134-L140)

The clear() implementation first deletes matching edges from memos_graph_edges and then deletes nodes from memos_graph_nodes. This test only asserts the node deletion, so a regression that removes the edge cleanup from clear() would go undetected — orphaned edges would remain in the database.

Consider adding an assertion to verify that DELETE FROM memos_graph_edges is also present in the executed SQL.

💡 Suggested Change

Before:

def test_clear_is_tenant_scoped(graph_db, cursor):
    cursor.fetchall.return_value = [("n1",)]
    graph_db.clear()
    sqls = _executed_sql(cursor)
    delete_nodes = [s for s in sqls if "DELETE FROM memos_graph_nodes" in s]
    assert delete_nodes
    assert all("user_name = %s" in s for s in delete_nodes)

After:

def test_clear_is_tenant_scoped(graph_db, cursor):
    cursor.fetchall.return_value = [("n1",)]
    graph_db.clear()
    sqls = _executed_sql(cursor)
    delete_nodes = [s for s in sqls if "DELETE FROM memos_graph_nodes" in s]
    assert delete_nodes
    assert all("user_name = %s" in s for s in delete_nodes)
    # Verify that associated edges are also removed (orphan cleanup).
    delete_edges = [s for s in sqls if "DELETE FROM memos_graph_edges" in s]
    assert delete_edges, "clear() must also delete edges to avoid orphaned rows"

6. tests/graph_dbs/test_oceanbase.py (L138-L140)

The clear() implementation first deletes matching edges from memos_graph_edges (to avoid orphaned rows), and then deletes nodes from memos_graph_nodes. This test only asserts on the node deletion, so a regression that drops the edge cleanup step would go undetected — orphaned edges would remain in the database.

Add an assertion to verify that DELETE FROM memos_graph_edges is also present in the executed SQL:

assert any("DELETE FROM memos_graph_edges" in s for s in sqls), (
    "clear() must also delete edges to avoid orphaned rows"
)
💡 Suggested Change

Before:

    delete_nodes = [s for s in sqls if "DELETE FROM memos_graph_nodes" in s]
    assert delete_nodes
    assert all("user_name = %s" in s for s in delete_nodes)

After:

    delete_nodes = [s for s in sqls if "DELETE FROM memos_graph_nodes" in s]
    assert delete_nodes
    assert all("user_name = %s" in s for s in delete_nodes)
    # Verify that associated edges are also deleted (orphan guard).
    assert any("DELETE FROM memos_graph_edges" in s for s in sqls), (
        "clear() must also delete edges to avoid orphaned rows"
    )

7. src/memos/graph_dbs/oceanbase.py (L194-L200)

If the second self._pool.acquire() returns another dead connection and ping raises again, that connection is never discarded — its slot stays counted in _created permanently. After enough such double-failures the pool starves: _created hits _maxconn but all slots are phantom-counted, so every subsequent acquire() blocks forever on _idle.get().

Apply the same ping-and-discard logic to the replacement connection:

💡 Suggested Change

Before:

        conn = self._pool.acquire()
        try:
            conn.ping(reconnect=True)
            return conn
        except Exception:
            self._pool.discard(conn)
            return self._pool.acquire()

After:

    def _acquire_live(self):
        while True:
            conn = self._pool.acquire()
            try:
                conn.ping(reconnect=True)
                return conn
            except Exception:
                self._pool.discard(conn)

8. src/memos/graph_dbs/oceanbase.py (L218-L232)

If conn.commit() or conn.rollback() raises (e.g. the TCP connection dropped mid-transaction), the connection is silently released back to the pool via self._pool.release(conn) in finally. The next borrower receives a connection that may still be in an open transaction or an error state, causing subtle data corruption or confusing errors.

Discard the connection instead of returning it to the pool when commit/rollback itself fails:

💡 Suggested Change

Before:

        try:
            cur = conn.cursor()
            conn.begin()
            try:
                yield cur
                conn.commit()
            except Exception:
                with suppress(Exception):
                    conn.rollback()
                raise
        finally:
            if cur is not None:
                with suppress(Exception):
                    cur.close()
            self._pool.release(conn)

After:

        committed = False
        try:
            cur = conn.cursor()
            conn.begin()
            try:
                yield cur
                conn.commit()
                committed = True
            except Exception:
                with suppress(Exception):
                    conn.rollback()
                raise
        finally:
            if cur is not None:
                with suppress(Exception):
                    cur.close()
            if committed:
                self._pool.release(conn)
            else:
                self._pool.discard(conn)

9. src/memos/graph_dbs/oceanbase.py (L936-L945)

vec is the string result of _vec_literal() (e.g. "[0.1,0.2,...]"). It is passed as a %s bound parameter twice. Whether OceanBase/pymysql can bind a plain Python string to a VECTOR parameter is driver- and server-version-specific and is not tested by the existing test suite. The vector literal is already safe to inline (it is constructed entirely from float() values with no user input), so inline it directly into the SQL to avoid the potential type-binding issue and to make the intent clear — consistent with how _vec_literal is used in add_node and update_node:

💡 Suggested Change

Before:

        rows = self._query(
            f"""
            SELECT id, 1 - cosine_distance(embedding, %s) AS score
            FROM {self.nodes_table}
            WHERE {where_clause}
            ORDER BY cosine_distance(embedding, %s)
            LIMIT %s
            """,
            (vec, *params, vec, top_k),
        )

After:

        rows = self._query(
            f"""
            SELECT id, 1 - cosine_distance(embedding, '{vec}') AS score
            FROM {self.nodes_table}
            WHERE {where_clause}
            ORDER BY cosine_distance(embedding, '{vec}')
            LIMIT %s
            """,
            (*params, top_k),
        )

10. src/memos/graph_dbs/oceanbase.py (L413-L421)

Each add_node call acquires its own connection and commits independently. A failure mid-batch leaves a partial set of nodes inserted with no rollback. Wrap the loop in a single _transaction so the batch is all-or-nothing, and use the cursor directly to avoid repeated pool acquire/release overhead:

💡 Suggested Change

Before:

    def add_nodes_batch(self, nodes: list[dict[str, Any]], user_name: str | None = None) -> None:
        """Batch add memory nodes."""
        for node in nodes:
            self.add_node(
                id=node["id"],
                memory=node["memory"],
                metadata=node.get("metadata", {}),
                user_name=user_name,
            )

After:

    def add_nodes_batch(self, nodes: list[dict[str, Any]], user_name: str | None = None) -> None:
        """Batch add memory nodes (atomic: all succeed or all roll back)."""
        if not nodes:
            return
        with self._transaction():
            for node in nodes:
                self.add_node(
                    id=node["id"],
                    memory=node["memory"],
                    metadata=node.get("metadata", {}),
                    user_name=user_name,
                )

11. src/memos/graph_dbs/oceanbase.py (L669)

Typo in the method name: prams should be params. The same typo appears in the log message [delete_node_by_prams] inside the method body. Since this is a new public method, rename it now before it becomes part of the external API.

💡 Suggested Change

Before:

    def delete_node_by_prams(

After:

    def delete_node_by_params(

12. src/memos/graph_dbs/oceanbase.py (L108)

datetime.utcnow() is deprecated since Python 3.12 and returns a naïve datetime with no timezone info, which can lead to silent mix-ups with timezone-aware datetimes. The same pattern is used in update_node() and _normalize_dt(). Replace with the timezone-aware equivalent throughout:

from datetime import datetime, timezone
now = datetime.now(timezone.utc).replace(tzinfo=None).isoformat()

🧹 Filtered 1 low-confidence OCR finding(s) before posting/fix-loop (duplicate: 1).

Generated by cloud-assistant via Open Code Review.

…ection handling

- Updated pyseekdb version constraints in pyproject.toml to restrict to <1.5.0.
- Increased default embedding dimension in APIConfig from 768 to 1024.
- Improved connection handling in OceanBaseGraphDB and OceanBaseVecDB to ensure better resource management and error handling.
- Added validation for table prefix length in OceanBaseGraphDB to prevent identifier overflow.
- Enhanced logging for empty password configurations in OceanBaseGraphDB.
- Updated tests to reflect changes in search behavior and connection management.
@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (31/31 executed). memos_python_core/changed-repo-python: 31/31. Duration: 5s

Branch: main

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 labels Jul 20, 2026
@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (31/31 executed). memos_python_core/changed-repo-python: 31/31. Duration: 6s

Branch: main

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Jul 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:api 云服务 / FastAPI / OpenAPI / MCP area:core MOS 编排层 / 框架底座 / 跨模块问题 area:database graph_db + vector_db | 图数据库与向量数据库 status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: support OceanBase for relational, vector, and graph storage

3 participants