From b0ed6fe0f3e5941ef90a08b3fa6cea3c1a7f6849 Mon Sep 17 00:00:00 2001
From: erie0210 <2j00923@gmail.com>
Date: Tue, 31 Oct 2023 22:00:52 +0900
Subject: [PATCH] feat: add sample stage

---
 .../core/aggregation/AggregationDsl.kt        | 10 ++++++
 .../core/aggregation/AggregationDslTest.kt    | 12 +++++++
 .../mongodb/repository/SampleRepository.kt    | 29 ++++++++++++++++
 .../repository/SampleRepositoryTest.kt        | 34 +++++++++++++++++++
 4 files changed, 85 insertions(+)
 create mode 100644 example/spring-data-mongodb/src/main/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepository.kt
 create mode 100644 example/spring-data-mongodb/src/test/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepositoryTest.kt

diff --git a/core/src/main/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDsl.kt b/core/src/main/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDsl.kt
index c6c3233a..f28d80d2 100644
--- a/core/src/main/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDsl.kt
+++ b/core/src/main/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDsl.kt
@@ -308,6 +308,16 @@ class AggregationDsl {
         operations += SetStageDsl().apply(configuration).get()
     }
 
+    /**
+     * Configures a stage that randomly selects the specified number of documents from the input documents to the pipeline.
+     *
+     * @param sampleSize Must not be less than zero.
+     * @see <a href="https:://www.mongodb.com/docs/manual/reference/operator/aggregation/sample/">$sample (aggregation)</a>
+     */
+    fun sample(sampleSize: Long) {
+        operations += Aggregation.sample(sampleSize)
+    }
+
     /**
      * Builds the [Aggregation] using the configured [AggregationOperation]s.
      *
diff --git a/core/src/test/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDslTest.kt b/core/src/test/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDslTest.kt
index 4dd23b69..a9acac99 100644
--- a/core/src/test/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDslTest.kt
+++ b/core/src/test/kotlin/com/github/inflab/spring/data/mongodb/core/aggregation/AggregationDslTest.kt
@@ -329,4 +329,16 @@ internal class AggregationDslTest : FreeSpec({
             ).toString()
         }
     }
+
+    "sample" {
+        // when
+        val aggregation = aggregation {
+            sample(3L)
+        }
+
+        // then
+        aggregation.toString() shouldBe Aggregation.newAggregation(
+            Aggregation.sample(3L),
+        ).toString()
+    }
 })
diff --git a/example/spring-data-mongodb/src/main/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepository.kt b/example/spring-data-mongodb/src/main/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepository.kt
new file mode 100644
index 00000000..15efe018
--- /dev/null
+++ b/example/spring-data-mongodb/src/main/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepository.kt
@@ -0,0 +1,29 @@
+package com.github.inflab.example.spring.data.mongodb.repository
+
+import com.github.inflab.spring.data.mongodb.core.aggregation.aggregation
+import org.springframework.data.annotation.Id
+import org.springframework.data.mongodb.core.MongoTemplate
+import org.springframework.data.mongodb.core.aggregate
+import org.springframework.data.mongodb.core.aggregation.AggregationResults
+import org.springframework.data.mongodb.core.mapping.Document
+import org.springframework.stereotype.Repository
+
+@Repository
+class SampleRepository(
+    private val mongoTemplate: MongoTemplate,
+) {
+
+    @Document("user")
+    data class User(@Id val id: String, val name: String, val q1: Boolean, val q2: Boolean)
+
+    /**
+     * @see <a href="https://www.mongodb.com/docs/manual/reference/operator/aggregation/sample/#example">Example</a>
+     */
+    fun example(): AggregationResults<User> {
+        val aggregation = aggregation {
+            sample(3)
+        }
+
+        return mongoTemplate.aggregate<User, User>(aggregation)
+    }
+}
diff --git a/example/spring-data-mongodb/src/test/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepositoryTest.kt b/example/spring-data-mongodb/src/test/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepositoryTest.kt
new file mode 100644
index 00000000..d0ba1aff
--- /dev/null
+++ b/example/spring-data-mongodb/src/test/kotlin/com/github/inflab/example/spring/data/mongodb/repository/SampleRepositoryTest.kt
@@ -0,0 +1,34 @@
+package com.github.inflab.example.spring.data.mongodb.repository
+
+import com.github.inflab.example.spring.data.mongodb.extension.makeMongoTemplate
+import com.github.inflab.example.spring.data.mongodb.repository.SampleRepository.User
+import io.kotest.core.spec.style.FreeSpec
+import io.kotest.matchers.shouldBe
+import org.springframework.data.mongodb.core.MongoTemplate
+
+internal class SampleRepositoryTest : FreeSpec({
+    val mongoTemplate: MongoTemplate = makeMongoTemplate()
+    val sampleRepository = SampleRepository(mongoTemplate)
+
+    beforeSpec {
+        val documents = listOf<User>(
+            User(id = "1", name = "dave123", q1 = true, q2 = true),
+            User(id = "2", name = "dave2", q1 = false, q2 = false),
+            User(id = "3", name = "ahn", q1 = true, q2 = true),
+            User(id = "4", name = "li", q1 = true, q2 = false),
+            User(id = "5", name = "annT", q1 = false, q2 = true),
+            User(id = "6", name = "li", q1 = true, q2 = true),
+            User(id = "7", name = "ty", q1 = false, q2 = true),
+        )
+
+        mongoTemplate.insertAll(documents)
+    }
+
+    "example" {
+        // when
+        val result = sampleRepository.example()
+
+        // then
+        result.mappedResults.toSet().size shouldBe 3
+    }
+})