From 37f26c412f06cf0435566456da0818cad3cf9ce6 Mon Sep 17 00:00:00 2001
From: He Wang <wanghechn@qq.com>
Date: Wed, 18 Jun 2025 14:00:59 +0800
Subject: [PATCH] add healthcheck to oceanbase container (#20989)

---
 .github/workflows/vdb-tests.yml               |  6 +--
 .../vdb/oceanbase/check_oceanbase_ready.py    | 49 -------------------
 docker/docker-compose-template.yaml           |  8 +++
 docker/docker-compose.yaml                    |  8 +++
 4 files changed, 18 insertions(+), 53 deletions(-)
 delete 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 512d14b2e..7d0a873eb 100644
--- a/.github/workflows/vdb-tests.yml
+++ b/.github/workflows/vdb-tests.yml
@@ -84,10 +84,8 @@ jobs:
             elasticsearch
             oceanbase
 
-      - name: Check VDB Ready (TiDB, Oceanbase)
-        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: Check VDB Ready (TiDB)
+        run: uv run --project api python api/tests/integration_tests/vdb/tidb_vector/check_tiflash_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
deleted file mode 100644
index 94a51292f..000000000
--- a/api/tests/integration_tests/vdb/oceanbase/check_oceanbase_ready.py
+++ /dev/null
@@ -1,49 +0,0 @@
-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()
diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml
index 55e1b5559..158ede88c 100644
--- a/docker/docker-compose-template.yaml
+++ b/docker/docker-compose-template.yaml
@@ -451,6 +451,14 @@ services:
       OB_CLUSTER_NAME: ${OCEANBASE_CLUSTER_NAME:-difyai}
       OB_SERVER_IP: 127.0.0.1
       MODE: mini
+    ports:
+      - "${OCEANBASE_VECTOR_PORT:-2881}:2881"
+    healthcheck:
+      test: [ 'CMD-SHELL', 'obclient -h127.0.0.1 -P2881 -uroot@test -p$${OB_TENANT_PASSWORD} -e "SELECT 1;"' ]
+      interval: 10s
+      retries: 30
+      start_period: 30s
+      timeout: 10s
 
   # Oracle vector database
   oracle:
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 2b98d098b..99aa87bcc 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -958,6 +958,14 @@ services:
       OB_CLUSTER_NAME: ${OCEANBASE_CLUSTER_NAME:-difyai}
       OB_SERVER_IP: 127.0.0.1
       MODE: mini
+    ports:
+      - "${OCEANBASE_VECTOR_PORT:-2881}:2881"
+    healthcheck:
+      test: [ 'CMD-SHELL', 'obclient -h127.0.0.1 -P2881 -uroot@test -p$${OB_TENANT_PASSWORD} -e "SELECT 1;"' ]
+      interval: 10s
+      retries: 30
+      start_period: 30s
+      timeout: 10s
 
   # Oracle vector database
   oracle: