Merge pull request #2046 from actiontech/issue-1191-2

feat:sql manage add new field endpoint
actiontech · Nov 17, 2023 · 8eeba6a · 8eeba6a
2 parents 50274ff + adde5d6
commit 8eeba6a
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 5 deletions.
diff --git a/sqle/api/controller/v1/audit_plan.go b/sqle/api/controller/v1/audit_plan.go
@@ -1006,6 +1006,7 @@ func convertToModelAuditPlanSQL(c echo.Context, auditPlan *model.AuditPlan, reqS
 			SQLContent:  reqSQL.LastReceiveText,
 			Info:        info,
 			Schema:      reqSQL.Schema,
+			Endpoint:    reqSQL.Endpoint,
 		})
 	}
 	return sqls, nil

diff --git a/sqle/model/audit_plan.go b/sqle/model/audit_plan.go
@@ -50,6 +50,8 @@ type AuditPlanSQLV2 struct {
 	SQLContent     string `json:"sql" gorm:"type:mediumtext;not null"`
 	Info           JSON   `gorm:"type:json"`
 	Schema         string `json:"schema" gorm:"type:varchar(512);not null"`
+	// 可能是URL或IP地址等
+	Endpoint string `json:"endpoint" gorm:"type:varchar(512)"`
 }
 
 type BlackListAuditPlanSQL struct {
@@ -280,10 +282,10 @@ ON DUPLICATE KEY UPDATE sql_content = VALUES(sql_content),
 func getBatchInsertRawSQL(auditPlanId uint, sqls []*AuditPlanSQLV2) (raw string, args []interface{}) {
 	pattern := make([]string, 0, len(sqls))
 	for _, sql := range sqls {
-		pattern = append(pattern, "(?, ?, ?, ?, ?, ?)")
-		args = append(args, auditPlanId, sql.GetFingerprintMD5(), sql.Fingerprint, sql.SQLContent, sql.Info, sql.Schema)
+		pattern = append(pattern, "(?, ?, ?, ?, ?, ?, ?)")
+		args = append(args, auditPlanId, sql.GetFingerprintMD5(), sql.Fingerprint, sql.SQLContent, sql.Info, sql.Schema, sql.Endpoint)
 	}
-	raw = fmt.Sprintf("INSERT INTO `audit_plan_sqls_v2` (`audit_plan_id`,`fingerprint_md5`, `fingerprint`, `sql_content`, `info`, `schema`) VALUES %s",
+	raw = fmt.Sprintf("INSERT INTO `audit_plan_sqls_v2` (`audit_plan_id`,`fingerprint_md5`, `fingerprint`, `sql_content`, `info`, `schema`, `endpoint`) VALUES %s",
 		strings.Join(pattern, ", "))
 	return
 }

diff --git a/sqle/model/audit_plan_test.go b/sqle/model/audit_plan_test.go
@@ -139,6 +139,7 @@ func TestStorage_OverrideAuditPlanSQLs(t *testing.T) {
 			SQLContent:  "select * from t1 where id = 1",
 			Info:        []byte(`{"counter": 1, "last_receive_timestamp": "mock time"}`),
 			Schema:      "test1",
+			Endpoint:    "10.186.1.2",
 		},
 	}
 
@@ -149,8 +150,8 @@ func TestStorage_OverrideAuditPlanSQLs(t *testing.T) {
 		WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectCommit()
 
-	mock.ExpectExec("INSERT INTO `audit_plan_sqls_v2` (`audit_plan_id`,`fingerprint_md5`, `fingerprint`, `sql_content`, `info`, `schema`) VALUES (?, ?, ?, ?, ?, ?);").
-		WithArgs(ap.ID, sqls[0].GetFingerprintMD5(), sqls[0].Fingerprint, sqls[0].SQLContent, sqls[0].Info, sqls[0].Schema).
+	mock.ExpectExec("INSERT INTO `audit_plan_sqls_v2` (`audit_plan_id`,`fingerprint_md5`, `fingerprint`, `sql_content`, `info`, `schema`, `endpoint`) VALUES (?, ?, ?, ?, ?, ?, ?);").
+		WithArgs(ap.ID, sqls[0].GetFingerprintMD5(), sqls[0].Fingerprint, sqls[0].SQLContent, sqls[0].Info, sqls[0].Schema, sqls[0].Endpoint).
 		WillReturnResult(sqlmock.NewResult(1, 1))
 
 	err = GetStorage().OverrideAuditPlanSQLs(ap.ID, sqls)

diff --git a/sqle/model/sql_manage.go b/sqle/model/sql_manage.go
@@ -17,6 +17,7 @@ type SqlManage struct {
 	LastReceiveTimestamp      *time.Time   `json:"last_receive_timestamp"`
 	InstanceName              string       `json:"instance_name"`
 	SchemaName                string       `json:"schema_name"`
+	Endpoint                  string       `json:"endpoint"`
 
 	Assignees []*User `gorm:"many2many:sql_manage_assignees;"`
 	Status    string  `json:"status" gorm:"default:\"unhandled\""`

diff --git a/sqle/server/auditplan/task.go b/sqle/server/auditplan/task.go
@@ -55,6 +55,7 @@ type SQL struct {
 	SQLContent  string
 	Fingerprint string
 	Schema      string
+	Endpoint    string
 	Info        map[string]interface{}
 }
 
@@ -277,6 +278,7 @@ func convertSQLsToModelSQLs(sqls []*SQL) []*model.AuditPlanSQLV2 {
 			SQLContent:  sql.SQLContent,
 			Schema:      sql.Schema,
 			Info:        data,
+			Endpoint:    sql.Endpoint,
 		}
 	}
 	return as