Skip to content

Commit

Permalink
feat: add sample stage
Browse files Browse the repository at this point in the history
  • Loading branch information
erie0210 committed Oct 31, 2023
1 parent d1ee223 commit b0ed6fe
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
})
Original file line number Diff line number Diff line change
@@ -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)
}
}
Original file line number Diff line number Diff line change
@@ -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
}
})

0 comments on commit b0ed6fe

Please sign in to comment.