From a13b697dfb4fe34b2bbbbbdee4c84b455846edfc Mon Sep 17 00:00:00 2001
From: Ben <182338249@qq.com>
Date: Sat, 27 Jan 2024 20:16:08 +0800
Subject: [PATCH] featrue: if id assigned,use it, otherwise use SnowFlakeId.

---
 .../common/entity/alerter/Alert.java          |  4 +++-
 .../common/entity/warehouse/History.java      |  4 +++-
 .../common/util/SnowFlakeIdGenerator.java     | 19 ++++++++++++++++++-
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/common/src/main/java/org/dromara/hertzbeat/common/entity/alerter/Alert.java b/common/src/main/java/org/dromara/hertzbeat/common/entity/alerter/Alert.java
index 490c4f66be6..055ecf22df1 100644
--- a/common/src/main/java/org/dromara/hertzbeat/common/entity/alerter/Alert.java
+++ b/common/src/main/java/org/dromara/hertzbeat/common/entity/alerter/Alert.java
@@ -23,6 +23,7 @@
 import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
+import org.hibernate.annotations.GenericGenerator;
 import org.hibernate.validator.constraints.Length;
 import org.springframework.data.annotation.CreatedBy;
 import org.springframework.data.annotation.CreatedDate;
@@ -55,7 +56,8 @@
 public class Alert {
 
     @Id
-    @GeneratedValue(strategy = GenerationType.IDENTITY)
+    @GeneratedValue(strategy = GenerationType.IDENTITY, generator = "myid")
+    @GenericGenerator(name = "myid", strategy = "org.dromara.hertzbeat.common.util.SnowFlakeIdGenerator")
     @Schema(title = "Alarm record entity primary key index ID",
             description = "告警记录实体主键索引ID",
             example = "87584674384", accessMode = READ_ONLY)
diff --git a/common/src/main/java/org/dromara/hertzbeat/common/entity/warehouse/History.java b/common/src/main/java/org/dromara/hertzbeat/common/entity/warehouse/History.java
index e1f3316ef84..3923b0478b3 100644
--- a/common/src/main/java/org/dromara/hertzbeat/common/entity/warehouse/History.java
+++ b/common/src/main/java/org/dromara/hertzbeat/common/entity/warehouse/History.java
@@ -5,6 +5,7 @@
 import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
+import org.hibernate.annotations.GenericGenerator;
 
 import javax.persistence.*;
 
@@ -30,7 +31,8 @@
 public class History {
 
     @Id
-    @GeneratedValue(strategy = GenerationType.IDENTITY)
+    @GeneratedValue(strategy = GenerationType.IDENTITY, generator = "myid")
+    @GenericGenerator(name = "myid", strategy = "org.dromara.hertzbeat.common.util.SnowFlakeIdGenerator")
     @Schema(description = "指标数据历史实体主键索引ID", example = "87584674384", accessMode = READ_ONLY)
     private Long id;
 
diff --git a/common/src/main/java/org/dromara/hertzbeat/common/util/SnowFlakeIdGenerator.java b/common/src/main/java/org/dromara/hertzbeat/common/util/SnowFlakeIdGenerator.java
index fdd48b411fd..109b1c618ad 100644
--- a/common/src/main/java/org/dromara/hertzbeat/common/util/SnowFlakeIdGenerator.java
+++ b/common/src/main/java/org/dromara/hertzbeat/common/util/SnowFlakeIdGenerator.java
@@ -17,11 +17,17 @@
 
 package org.dromara.hertzbeat.common.util;
 
+import org.hibernate.HibernateException;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.id.IdentityGenerator;
+
+import java.io.Serializable;
+
 /**
  * Snowflake Algorithm Generator Tool
  * @author tomsun28
  */
-public class SnowFlakeIdGenerator {
+public class SnowFlakeIdGenerator extends IdentityGenerator {
 
     private static final SnowFlakeIdWorker ID_WORKER;
 
@@ -29,6 +35,17 @@ public class SnowFlakeIdGenerator {
         ID_WORKER = new SnowFlakeIdWorker();
     }
 
+    @Override
+    public Serializable generate(SharedSessionContractImplementor s, Object obj) throws HibernateException {
+        Serializable id = s.getEntityPersister(null, obj).getClassMetadata().getIdentifier(obj, s);
+
+        if (id != null && Long.valueOf(id.toString()) > 0) {
+            return id;
+        } else {
+            return SnowFlakeIdGenerator.generateId();
+        }
+    }
+
     public static long generateId() {
         return ID_WORKER.nextId();
     }