Skip to content

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

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

feat: add OceanBase/seekdb vector and graph storage providers#2130
Evenss wants to merge 3 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: abad4fed10e4fccd
Base: main
Head: main

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

⚠️ 2 warning(s) occurred during review.


1. pyproject.toml (L117)

The upper bound <1.5.0 is a patch-level pin that will block any 1.5.x release of pyseekdb. Other optional extras in this file (e.g. neo4j, pymilvus, fastapi) use a minor-version upper bound (<6.0.0, <3.0.0) to stay compatible with bug-fix releases while still guarding against breaking major/minor bumps. Consider widening to >=1.4.0,<2.0.0 (or at least <1.6.0) unless there is a known breaking change in 1.5.0 that justifies the tighter bound.

💡 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/api/config.py (L926)

The OceanBase config re-uses the same EMBEDDING_DIMENSION environment variable as get_postgres_config, but with a different default value (1024 here vs. 384 in postgres). When EMBEDDING_DIMENSION is not explicitly set in the environment, the effective default depends on which method is called — this can silently produce the wrong dimension if an operator switches backends without updating the env var, leading to schema/index incompatibility at runtime.

Consider using a dedicated variable (e.g., OCEANBASE_EMBEDDING_DIMENSION) that falls back to the generic one:

"embedding_dimension": int(os.getenv("OCEANBASE_EMBEDDING_DIMENSION") or os.getenv("EMBEDDING_DIMENSION", "1024")),

3. src/memos/api/config.py (L920)

Defaulting to "root" is a privileged database account with unrestricted access. If OCEANBASE_USER is not set (e.g., during development or in a misconfigured deployment), the application connects as root, violating the principle of least privilege and increasing the blast radius of any SQL injection or misconfiguration.

Consider a less-privileged default (e.g., "memos") and document that OCEANBASE_USER must be set explicitly in production.


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

Bug: validator silently bypassed when vector_dimension is omitted.

In Pydantic v2, @field_validator does not run against a field's default value unless validate_default=True is specified. Because vector_dimension is inherited with default=None, creating OceanBaseVecDBConfig(host=..., collection_name=...) without providing vector_dimension will skip this validator and produce a config where vector_dimension is None. The error is then deferred until HNSWConfiguration(dimension=None, ...) is called at runtime.

Two clean fixes:

Option A – mark the validator to run on defaults (minimal change):

@field_validator("vector_dimension", mode="before")
@classmethod
def validate_vector_dimension(cls, value: int | None) -> int:

or

@field_validator("vector_dimension", validate_default=True)

Option B – override the field to be required (most explicit):

vector_dimension: int = Field(..., description="Dimension of the vectors")

This makes the missing-dimension error a clear Pydantic validation failure at config construction time instead of a cryptic runtime error inside create_collection().

💡 Suggested Change

Before:

    @field_validator("vector_dimension")
    @classmethod
    def validate_vector_dimension(cls, value: int | None) -> int:
        """OceanBase requires a concrete vector dimension to build the HNSW index."""
        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 value

After:

    @field_validator("vector_dimension", mode="before")
    @classmethod
    def validate_vector_dimension(cls, value: int | None) -> int:
        """OceanBase requires a concrete vector dimension to build the HNSW index."""
        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 value

5. tests/graph_dbs/test_oceanbase.py (L76)

The implementation's _parse_row injects both created_at (row[3]) and updated_at (row[4]) into the returned metadata dict. The test asserts created_at but never checks updated_at, leaving a regression gap. Adding an assertion for updated_at would make the contract explicit and catch any future key-rename or injection bug for the second timestamp field.

💡 Suggested Change

Before:

    assert node["metadata"]["created_at"] == "2024-01-01 00:00:00"

After:

    assert node["metadata"]["created_at"] == "2024-01-01 00:00:00"
    assert node["metadata"]["updated_at"] == "2024-01-02 00:00:00"

6. tests/vec_dbs/test_oceanbase.py (L87-L89)

The comment says "nested included" but the test data is a flat dict with no nested structure. This is misleading — a reader would expect a nested payload (e.g., {"memory": "hi", "metadata": {"status": "activated"}}) to demonstrate nesting support. Either add a nested key to the test data to actually exercise the claim, or remove the parenthetical.


7. src/memos/graph_dbs/oceanbase.py (L202-L210)

Bug: connection leaked on _acquire_live failure path in _borrow and _transaction.

If _acquire_live() raises (e.g. a second acquire() fails because the pool is exhausted or the server is down), conn is never assigned, so the finally block calls self._pool.release(conn) with an unbound name — raising a NameError instead of the original exception. The fix is to initialize conn = None before the call and guard the release accordingly.

Same pattern applies in _transaction.

💡 Suggested Change

Before:

    @contextmanager
    def _borrow(self):
        """Borrow a single connection for one auto-committed operation."""
        self._ensure_open()
        conn = self._acquire_live()
        try:
            yield conn
        finally:
            self._pool.release(conn)

After:

    @contextmanager
    def _borrow(self):
        """Borrow a single connection for one auto-committed operation."""
        self._ensure_open()
        conn = None
        conn = self._acquire_live()
        try:
            yield conn
        finally:
            if conn is not None:
                self._pool.release(conn)

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

Bug: connection leaked when _transaction raises before cur = conn.cursor().

If conn.begin() or conn.cursor() raises, the finally block closes cur (correctly guarded), but then calls self._pool.release(conn). However, conn is already bound so the release itself is fine — the real gap is that rollback is only attempted in the inner except, not in the path where conn.begin() throws before the inner try is entered. In that case the connection is returned to the pool in a potentially uncommitted state from a prior (auto-committed) operation. More critically: if conn was never obtained (parallel to the _borrow issue above, though here it IS obtained before the try), this is safe, but conn.begin() raising before the inner try means no rollback is issued. Add an outer except or restructure so rollback is always attempted on any exception path.

💡 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:

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

9. src/memos/graph_dbs/oceanbase.py (L71)

Bug: acquire() can block forever (deadlock) when all connections are in-flight.

When self._created == self._maxconn and the idle queue is empty, acquire() falls through to self._idle.get() — a blocking call with no timeout. If a thread holding the last connection is waiting for another connection (or the caller forgot to release), this deadlocks silently. Consider using self._idle.get(timeout=...) and raising a GraphDBError on timeout so callers get a clear error rather than hanging indefinitely.

💡 Suggested Change

Before:

        return self._idle.get()

After:

        try:
            return self._idle.get(timeout=30)
        except queue.Empty:
            raise GraphDBError(
                f"Connection pool exhausted (maxconn={self._maxconn}): "
                "timed out waiting for a free connection"
            )

10. src/memos/graph_dbs/oceanbase.py (L89-L93)

Bug: release() leaks a pool slot when put_nowait raises queue.Full.

When the idle queue is full, the connection is closed but self._created is never decremented. Subsequent acquire() calls will see self._created == self._maxconn and block on self._idle.get() forever, even though the pool has capacity. Use self._pool.discard(conn) instead of conn.close() so the counter stays accurate.

💡 Suggested Change

Before:

        try:
            self._idle.put_nowait(conn)
        except queue.Full:
            with suppress(Exception):
                conn.close()

After:

        try:
            self._idle.put_nowait(conn)
        except queue.Full:
            self.discard(conn)

11. src/memos/graph_dbs/oceanbase.py (L154-L155)

Bug: SQL injection via unsanitized table_prefix in DDL.

self.nodes_table and self.edges_table are interpolated directly into f-string DDL/DML throughout the file (e.g. CREATE TABLE IF NOT EXISTS {self.nodes_table}). Although OceanBaseGraphDBConfig.validate_table_prefix restricts table_prefix to [A-Za-z_][A-Za-z0-9_]*, the constructed table names ({prefix}_nodes, {prefix}_edges) are never back-quoted / escaped before being embedded in SQL. If the validator is bypassed or the config is constructed without Pydantic validation, the table names reach the database raw. A belt-and-suspenders fix is to apply the same regex check on the final table name strings at construction time in __init__, or wrap them in backticks in every SQL statement.

💡 Suggested Change

Before:

        self.nodes_table = f"{config.table_prefix}_nodes"
        self.edges_table = f"{config.table_prefix}_edges"

After:

        _safe = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*$')
        nodes_table = f"{config.table_prefix}_nodes"
        edges_table = f"{config.table_prefix}_edges"
        if not _safe.match(nodes_table) or not _safe.match(edges_table):
            raise ValueError(f"Derived table names are not safe identifiers: {nodes_table}, {edges_table}")
        self.nodes_table = f"`{nodes_table}`"
        self.edges_table = f"`{edges_table}`"

12. src/memos/graph_dbs/oceanbase.py (L937-L944)

Bug: vector parameter positional mismatch in search_by_embedding.

The SQL is:

SELECT id, 1 - cosine_distance(embedding, %s) AS score
FROM ...
WHERE {where_clause}
ORDER BY cosine_distance(embedding, %s)
LIMIT %s

The params tuple passed is (vec, *params, vec, top_k). params grows dynamically from the conditions list (user_name, scope, status, search_filter fields). The first %s in SELECT is bound to vec correctly, but the second %s (in ORDER BY) is meant to be vec and top_k is the third — this order is only correct if the WHERE-clause params are placed between the two vec values, which they are. However the SQL text has %s for the SELECT clause before the WHERE clause %s placeholders, but the driver processes all %s left-to-right across the entire statement. So the actual binding order is:

  1. SELECT %svec
  2. WHERE %s... → *params
  3. ORDER BY %svec
  4. LIMIT %stop_k

This is actually correct as written. However the construction is fragile — any reordering of clauses or addition of a WHERE param would silently corrupt results. Consider binding the vector once via a subquery or a user variable to make it impossible to get out of sync.


🧹 Filtered 1 low-confidence OCR finding(s) before posting/fix-loop (existing_code_mismatch: 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
@wustzdy

wustzdy commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

新增的db不符合
image
木项目的代码结构

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

Evenss commented Jul 27, 2026

Copy link
Copy Markdown
Author

@wustzdy 感谢 review!想跟您确认一下「不符合本项目代码结构」具体指的是哪方面,以便我针对性调整:

本次改动的文件布局是参照现有 polardb / postgres 后端的接入方式来做的:

  • 实现:src/memos/graph_dbs/oceanbase.pysrc/memos/vec_dbs/oceanbase.py(与 polardb.pypostgres.pymilvus.py一样平铺在各自目录下)
  • 配置:src/memos/configs/graph_db.py / vec_db.py 中新增 pydantic 配置并注册到 ConfigFactory
  • 工厂:graph_dbs/factory.py / vec_dbs/factory.py 注册 oceanbase / seekdb 别名
  • 依赖:pyseekdb 放在可选 extras ob-mem 中,import 有 try/except 保护,不影响核心依赖
  • 测试:tests/graph_dbs/tests/vec_dbs/ 下,与源码目录一一对应

如果您指的是其他结构问题,麻烦指明一下期望的结构,我尽快调整,谢谢!

@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: 8s [advisory, non-gating] AI-generated tests on branch test/auto-gen-abad4fed10e4fccd-20260727164640: 181/186 passed, 5 failed — these do NOT affect the PR verdict; review the branch manually.

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 27, 2026
@Evenss

Evenss commented Jul 30, 2026

Copy link
Copy Markdown
Author

@wustzdy PTAL

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

4 participants