From 00ca18ee44626318468a4b5618168ad9dc616d1d Mon Sep 17 00:00:00 2001
From: Bowen Liang <liangbowen@gf.com.cn>
Date: Fri, 27 Jun 2025 16:00:01 +0800
Subject: [PATCH] restore ob readiness detection

---
 .github/workflows/vdb-tests.yml               |  4 +-
 .../vdb/oceanbase/check_oceanbase_ready.py    | 49 +++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 api/tests/integration_tests/vdb/oceanbase/check_oceanbase_ready.py

diff --git a/.github/workflows/vdb-tests.yml b/.github/workflows/vdb-tests.yml
index 912267094b..6185743ece 100644
--- a/.github/workflows/vdb-tests.yml
+++ b/.github/workflows/vdb-tests.yml
@@ -91,7 +91,9 @@ jobs:
           cp api/tests/integration_tests/.env.example api/tests/integration_tests/.env
 
       - name: Check VDB Ready (TiDB)
-        run: uv run --project api python api/tests/integration_tests/vdb/tidb_vector/check_tiflash_ready.py
+        run: |
+          uv run --project api python api/tests/integration_tests/vdb/tidb_vector/check_tiflash_ready.py
+          uv run --project api python api/tests/integration_tests/vdb/oceanbase/check_oceanbase_ready.py
 
       - name: Test Vector Stores
         run: uv run --project api bash dev/pytest/pytest_vdb.sh
diff --git a/api/tests/integration_tests/vdb/oceanbase/check_oceanbase_ready.py b/api/tests/integration_tests/vdb/oceanbase/check_oceanbase_ready.py
new file mode 100644
index 0000000000..94a51292ff
--- /dev/null
+++ b/api/tests/integration_tests/vdb/oceanbase/check_oceanbase_ready.py
@@ -0,0 +1,49 @@
+import time
+
+import pymysql
+
+
+def check_oceanbase_ready() -> bool:
+    try:
+        connection = pymysql.connect(
+            host="localhost",
+            port=2881,
+            user="root",
+            password="difyai123456",
+        )
+        affected_rows = connection.query("SELECT 1")
+        return affected_rows == 1
+    except Exception as e:
+        print(f"Oceanbase is not ready. Exception: {e}")
+        return False
+    finally:
+        if connection:
+            connection.close()
+
+
+def main():
+    max_attempts = 50
+    retry_interval_seconds = 2
+    is_oceanbase_ready = False
+    for attempt in range(max_attempts):
+        try:
+            is_oceanbase_ready = check_oceanbase_ready()
+        except Exception as e:
+            print(f"Oceanbase is not ready. Exception: {e}")
+            is_oceanbase_ready = False
+
+        if is_oceanbase_ready:
+            break
+        else:
+            print(f"Attempt {attempt + 1} failed, retry in {retry_interval_seconds} seconds...")
+            time.sleep(retry_interval_seconds)
+
+    if is_oceanbase_ready:
+        print("Oceanbase is ready.")
+    else:
+        print(f"Oceanbase is not ready after {max_attempts} attempting checks.")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()