From 323ee1f7270af968fc3af4f3e1de619d24514e8d Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Fri, 14 Jul 2023 10:04:44 +0800
Subject: [PATCH] =?UTF-8?q?=E6=89=8B=E5=8A=A8=E5=90=8C=E6=AD=A5=E9=9B=86?=
 =?UTF-8?q?=E7=BE=A4=E4=BB=BB=E5=8A=A1=E5=90=8E=E6=8A=8A=E6=89=80=E6=9C=89?=
 =?UTF-8?q?=E7=9B=B8=E5=85=B3=E4=BB=BB=E5=8A=A1=E6=A0=87=E8=AE=B0=E4=B8=BA?=
 =?UTF-8?q?=E5=B7=B2=E5=AE=8C=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/db/models/dns/dns_task_dao.go      | 21 +++++++++++++++++++++
 internal/db/models/dns/dns_task_dao_test.go | 15 +++++++++++++--
 internal/rpc/services/service_dns_domain.go | 19 ++++++++++++++++++-
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/internal/db/models/dns/dns_task_dao.go b/internal/db/models/dns/dns_task_dao.go
index 4b8b34a8..d7c94034 100644
--- a/internal/db/models/dns/dns_task_dao.go
+++ b/internal/db/models/dns/dns_task_dao.go
@@ -201,6 +201,27 @@ func (this *DNSTaskDAO) UpdateDNSTaskDone(tx *dbs.Tx, taskId int64, taskVersion
 	return this.Save(tx, op)
 }
 
+// GenerateVersion 生成最新的版本号
+func (this *DNSTaskDAO) GenerateVersion() int64 {
+	return time.Now().UnixNano()
+}
+
+// UpdateClusterDNSTasksDone 设置所有集群任务完成
+func (this *DNSTaskDAO) UpdateClusterDNSTasksDone(tx *dbs.Tx, clusterId int64, maxVersion int64) error {
+	if clusterId <= 0 || maxVersion <= 0 {
+		return nil
+	}
+
+	return this.Query(tx).
+		Attr("clusterId", clusterId).
+		Attr("isOk", false).
+		Lte("version", maxVersion).
+		Set("isDone", true).
+		Set("isOk", true).
+		Set("error", "").
+		UpdateQuickly()
+}
+
 // DeleteDNSTasksWithClusterId 删除集群相关任务
 func (this *DNSTaskDAO) DeleteDNSTasksWithClusterId(tx *dbs.Tx, clusterId int64) error {
 	if clusterId <= 0 {
diff --git a/internal/db/models/dns/dns_task_dao_test.go b/internal/db/models/dns/dns_task_dao_test.go
index 17449080..5c8fcdee 100644
--- a/internal/db/models/dns/dns_task_dao_test.go
+++ b/internal/db/models/dns/dns_task_dao_test.go
@@ -1,17 +1,28 @@
-package dns
+package dns_test
 
 import (
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models/dns"
 	_ "github.com/go-sql-driver/mysql"
 	_ "github.com/iwind/TeaGo/bootstrap"
 	"github.com/iwind/TeaGo/dbs"
 	"testing"
+	"time"
 )
 
 func TestDNSTaskDAO_CreateDNSTask(t *testing.T) {
 	dbs.NotifyReady()
-	err := SharedDNSTaskDAO.CreateDNSTask(nil, 1, 2, 3, 0, "cdn", "taskType")
+	err := dns.SharedDNSTaskDAO.CreateDNSTask(nil, 1, 2, 3, 0, "cdn", "taskType")
 	if err != nil {
 		t.Fatal(err)
 	}
 	t.Log("ok")
 }
+
+func TestDNSTaskDAO_UpdateClusterDNSTasksDone(t *testing.T) {
+	var dao = dns.NewDNSTaskDAO()
+	var tx *dbs.Tx
+	err := dao.UpdateClusterDNSTasksDone(tx, 46, time.Now().UnixNano())
+	if err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/internal/rpc/services/service_dns_domain.go b/internal/rpc/services/service_dns_domain.go
index 3161d0bd..416b8084 100644
--- a/internal/rpc/services/service_dns_domain.go
+++ b/internal/rpc/services/service_dns_domain.go
@@ -303,7 +303,24 @@ func (this *DNSDomainService) SyncDNSDomainData(ctx context.Context, req *pb.Syn
 	if err != nil {
 		return nil, err
 	}
-	return this.syncClusterDNS(req)
+
+	var latestVersion = dns.SharedDNSTaskDAO.GenerateVersion()
+
+	resp, err := this.syncClusterDNS(req)
+	if err != nil {
+		return resp, err
+	}
+
+	// 标记集群所有任务已完成
+	if req.NodeClusterId > 0 && resp != nil && resp.IsOk {
+		var tx = this.NullTx()
+		err = dns.SharedDNSTaskDAO.UpdateClusterDNSTasksDone(tx, req.NodeClusterId, latestVersion)
+		if err != nil {
+			return resp, err
+		}
+	}
+
+	return resp, err
 }
 
 // FindAllDNSDomainRoutes 查看支持的线路