From 063d93b9b32d51a84021dd95406d1ae405078cbe Mon Sep 17 00:00:00 2001
From: BugsGuru <lwqlwq888@gmail.com>
Date: Thu, 4 Dec 2025 16:00:56 +0800
Subject: [PATCH] feat: enhance DirectAuditParams with InstanceName and
 SchemaName

---
 internal/dms/biz/cloudbeaver.go     | 13 +++++++++----
 internal/pkg/cloudbeaver/graphql.go |  4 ++++
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/internal/dms/biz/cloudbeaver.go b/internal/dms/biz/cloudbeaver.go
index 28f1b232..8d6b4ccb 100644
--- a/internal/dms/biz/cloudbeaver.go
+++ b/internal/dms/biz/cloudbeaver.go
@@ -486,12 +486,18 @@ func (cu *CloudbeaverUsecase) GraphQLDistributor() echo.MiddlewareFunc {
 						return err
 					}
 
+					execParams, err := cu.getWorkflowExecParams(c, params)
+					if err != nil {
+						return err
+					}
 					// 构建直接审计请求参数
 					directAuditReq := cloudbeaver.DirectAuditParams{
 						AuditSQLReq: cloudbeaver.AuditSQLReq{
 							InstanceType:     dbService.DBType,
 							ProjectId:        dbService.ProjectUID,
 							RuleTemplateName: dbService.SQLEConfig.SQLQueryConfig.RuleTemplateName,
+							InstanceName:     dbService.Name,
+							SchemaName:       execParams.instanceSchema,
 						},
 						SQLEAddr:                         fmt.Sprintf("%s/v2/sql_audit", sqleUrl),
 						AllowQueryWhenLessThanAuditLevel: dbService.GetAllowQueryWhenLessThanAuditLevel(),
@@ -1941,13 +1947,12 @@ func (cu *CloudbeaverUsecase) getContextSchema(c echo.Context, connectionId, con
 	}
 
 	contextInfo := res.Contexts[0]
-	if contextInfo.DefaultSchema != nil && *contextInfo.DefaultSchema != "" {
-		return *contextInfo.DefaultSchema, nil
-	}
-
 	if contextInfo.DefaultCatalog != nil && *contextInfo.DefaultCatalog != "" {
 		return *contextInfo.DefaultCatalog, nil
 	}
+	if contextInfo.DefaultSchema != nil && *contextInfo.DefaultSchema != "" {
+		return *contextInfo.DefaultSchema, nil
+	}
 
 	return "", nil
 }
diff --git a/internal/pkg/cloudbeaver/graphql.go b/internal/pkg/cloudbeaver/graphql.go
index 7ef3876c..af2dd63a 100644
--- a/internal/pkg/cloudbeaver/graphql.go
+++ b/internal/pkg/cloudbeaver/graphql.go
@@ -123,6 +123,8 @@ type AuditSQLReq struct {
 	SQLType          string `json:"sql_type" form:"sql_type" example:"sql" enums:"sql,mybatis," valid:"omitempty,oneof=sql mybatis"`
 	ProjectId        string `json:"project_id" form:"project_id" example:"700300" valid:"required"`
 	RuleTemplateName string `json:"rule_template_name" form:"rule_template_name" example:"default" valid:"required"`
+	InstanceName     string `json:"instance_name" form:"instance_name" example:"instance1"`
+	SchemaName       string `json:"schema_name" form:"schema_name" example:"schema1"`
 }
 
 type DirectAuditParams struct {
@@ -174,6 +176,8 @@ func (r *MutationResolverImpl) AuditSQL(ctx context.Context, sql string, connect
 		SQLType:          "sql",
 		ProjectId:        directAuditParams.ProjectId,
 		RuleTemplateName: directAuditParams.RuleTemplateName,
+		InstanceName:     directAuditParams.InstanceName,
+		SchemaName:       directAuditParams.SchemaName,
 	}
 
 	reply := &auditSQLReply{}