From fc44e53d9a6c0d93a9f0a0e0d1b724089ce7886a Mon Sep 17 00:00:00 2001
From: Lynwee <linwei.hou@merico.dev>
Date: Tue, 17 Oct 2023 11:45:33 +0800
Subject: [PATCH] feat(zentao): check db url when testing connections (#6258)

* feat(zentao): check db url when testing connections

* fix(zentao): remove test codes

* fix(zentao): remove test codes
---
 backend/core/runner/db.go                     | 35 ++++++++++++++-----
 config-ui/src/api/connection/index.ts         |  2 +-
 config-ui/src/api/connection/types.ts         |  2 ++
 .../components/connection-form/index.tsx      |  1 +
 config-ui/src/store/connections/context.tsx   |  3 ++
 config-ui/src/store/connections/types.ts      |  1 +
 6 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/backend/core/runner/db.go b/backend/core/runner/db.go
index 72febf5046a..af0a3386c4c 100644
--- a/backend/core/runner/db.go
+++ b/backend/core/runner/db.go
@@ -160,15 +160,32 @@ func getDbConnection(dbUrl string, conf *gorm.Config) (*gorm.DB, error) {
 }
 
 func CheckDbConnection(dbUrl string, d time.Duration) errors.Error {
-	db, err := getDbConnection(dbUrl, &gorm.Config{})
-	if err != nil {
-		return errors.Convert(err)
-	}
 	ctx := context.Background()
-	if d > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(context.Background(), d)
-		defer cancel()
+
+	result := make(chan errors.Error, 1)
+	done := make(chan struct{}, 1)
+	go func() {
+		db, err := getDbConnection(dbUrl, &gorm.Config{})
+		if err != nil {
+			result <- errors.Convert(err)
+		}
+		if d > 0 {
+			var cancel context.CancelFunc
+			ctx, cancel = context.WithTimeout(context.Background(), d)
+			defer cancel()
+		}
+		if err := db.WithContext(ctx).Exec("SELECT 1").Error; err != nil {
+			done <- struct{}{}
+		} else {
+			result <- errors.Convert(err)
+		}
+	}()
+	select {
+	case <-time.After(d):
+		return errors.Default.New("timeout")
+	case <-done:
+		return nil
+	case err := <-result:
+		return err
 	}
-	return errors.Convert(db.WithContext(ctx).Exec("SELECT 1").Error)
 }
diff --git a/config-ui/src/api/connection/index.ts b/config-ui/src/api/connection/index.ts
index 138ea6b2e1f..c142d2879b9 100644
--- a/config-ui/src/api/connection/index.ts
+++ b/config-ui/src/api/connection/index.ts
@@ -41,6 +41,6 @@ export const test = (
   plugin: string,
   payload: Pick<
     T.ConnectionForm,
-    'endpoint' | 'authMethod' | 'username' | 'password' | 'token' | 'appId' | 'secretKey' | 'proxy'
+    'endpoint' | 'authMethod' | 'username' | 'password' | 'token' | 'appId' | 'secretKey' | 'proxy' | 'dbUrl'
   >,
 ): Promise<T.ConnectionTest> => request(`/plugins/${plugin}/test`, { method: 'post', data: payload });
diff --git a/config-ui/src/api/connection/types.ts b/config-ui/src/api/connection/types.ts
index 933133e2d37..a61f954a796 100644
--- a/config-ui/src/api/connection/types.ts
+++ b/config-ui/src/api/connection/types.ts
@@ -26,6 +26,7 @@ export type Connection = {
   password?: string;
   proxy: string;
   apiKey?: string;
+  dbUrl?: string;
 };
 
 export type ConnectionForm = {
@@ -40,6 +41,7 @@ export type ConnectionForm = {
   enableGraphql?: boolean;
   proxy: string;
   rateLimitPerHour?: number;
+  dbUrl?: string;
 };
 
 export type ConnectionTest = {
diff --git a/config-ui/src/plugins/components/connection-form/index.tsx b/config-ui/src/plugins/components/connection-form/index.tsx
index e64dd6ba4a7..4ef0a9573b9 100644
--- a/config-ui/src/plugins/components/connection-form/index.tsx
+++ b/config-ui/src/plugins/components/connection-form/index.tsx
@@ -73,6 +73,7 @@ export const ConnectionForm = ({ plugin, connectionId, onSuccess }: Props) => {
             'secretKey',
             'tenantId',
             'tenantType',
+            "dbUrl",
           ]),
         ),
       {
diff --git a/config-ui/src/store/connections/context.tsx b/config-ui/src/store/connections/context.tsx
index a41b238bcf9..42e213dea7f 100644
--- a/config-ui/src/store/connections/context.tsx
+++ b/config-ui/src/store/connections/context.tsx
@@ -72,6 +72,7 @@ export const ConnectionContextProvider = ({ children, ...props }: Props) => {
     authMethod,
     secretKey,
     appId,
+    dbUrl,
   }: ConnectionItemType) => {
     try {
       const res = await API.connection.test(plugin, {
@@ -83,6 +84,7 @@ export const ConnectionContextProvider = ({ children, ...props }: Props) => {
         authMethod,
         secretKey,
         appId,
+        dbUrl,
       });
       return res.success ? ConnectionStatusEnum.ONLINE : ConnectionStatusEnum.OFFLINE;
     } catch {
@@ -109,6 +111,7 @@ export const ConnectionContextProvider = ({ children, ...props }: Props) => {
       authMethod: it.authMethod,
       secretKey: it.secretKey,
       appId: it.appId,
+      dbUrl: it.dbUrl,
     }));
   };
 
diff --git a/config-ui/src/store/connections/types.ts b/config-ui/src/store/connections/types.ts
index 68847166f87..07475889a6d 100644
--- a/config-ui/src/store/connections/types.ts
+++ b/config-ui/src/store/connections/types.ts
@@ -41,4 +41,5 @@ export type ConnectionItemType = {
   authMethod?: string;
   appId?: string;
   secretKey?: string;
+  dbUrl?: string;
 };