From 4f2870cdd81d7d5e668c976662876ea4c3157743 Mon Sep 17 00:00:00 2001
From: Tomasz <tomasz.cichoszewski@h2o.ai>
Date: Mon, 22 Jan 2024 12:43:00 +0100
Subject: [PATCH] Mapped forceGuaranteedRollup parameter

---
 common_spec_types.go | 1 +
 task_types.go        | 8 ++++++++
 task_types_test.go   | 2 ++
 3 files changed, 11 insertions(+)

diff --git a/common_spec_types.go b/common_spec_types.go
index a6e6e7b..9bada0f 100644
--- a/common_spec_types.go
+++ b/common_spec_types.go
@@ -38,6 +38,7 @@ type TuningConfig struct {
 	MaxRowsPerSegment                int        `json:"maxRowsPerSegment,omitempty"`
 	MaxRowsInMemory                  int        `json:"maxRowsInMemory,omitempty"`
 	IndexSpecForIntermediatePersists *IndexSpec `json:"indexSpecForIntermediatePersists,omitempty"`
+	ForceGuaranteedRollup            bool       `json:"forceGuaranteedRollup,omitempty"`
 }
 
 // Metric is a Druid aggregator that is applied at ingestion time.
diff --git a/task_types.go b/task_types.go
index 8e7d9be..19d8b34 100644
--- a/task_types.go
+++ b/task_types.go
@@ -224,6 +224,14 @@ func SetTaskGranularitySpec(segmentGranularity string, queryGranularity *QueryGr
 	}
 }
 
+// SetForceGuaranteedRollup sets guaranteed rollup setting for ingestion spec.
+// https://druid.apache.org/docs/latest/ingestion/rollup#perfect-rollup-vs-best-effort-rollup
+func SetForceGuaranteedRollup(rollup bool) TaskIngestionSpecOptions {
+	return func(spec *TaskIngestionSpec) {
+		spec.Spec.TuningConfig.ForceGuaranteedRollup = rollup
+	}
+}
+
 // NewTaskIngestionSpec returns a default TaskIngestionSpec and applies any options passed to it.
 func NewTaskIngestionSpec(options ...TaskIngestionSpecOptions) *TaskIngestionSpec {
 	spec := defaultTaskIngestionSpec()
diff --git a/task_types_test.go b/task_types_test.go
index 2735d7b..b7868d4 100644
--- a/task_types_test.go
+++ b/task_types_test.go
@@ -26,6 +26,7 @@ func TestTaskIngestionSpec(t *testing.T) {
 				SetTaskSchemaDiscovery(false),
 				SetTaskTimestampColumn("__time"),
 				SetTaskGranularitySpec("DAY", &QueryGranularitySpec{"none"}, true),
+				SetForceGuaranteedRollup(true),
 			},
 			expected: func() *TaskIngestionSpec {
 				out := NewTaskIngestionSpec()
@@ -41,6 +42,7 @@ func TestTaskIngestionSpec(t *testing.T) {
 				out.Spec.DataSchema.GranularitySpec.SegmentGranularity = "DAY"
 				out.Spec.DataSchema.GranularitySpec.QueryGranularity = &QueryGranularitySpec{"none"}
 				out.Spec.DataSchema.GranularitySpec.Rollup = true
+				out.Spec.TuningConfig.ForceGuaranteedRollup = true
 				return out
 			}(),
 		},