Skip to content

Commit

Permalink
✨ Source MongoDB Internal POC: Generate Test Data (#29049)
Browse files Browse the repository at this point in the history
* Add script to generate test data

* Fix prose

* Update credentials example

* PR feedback
  • Loading branch information
jdpgrailsdev authored Aug 3, 2023
1 parent 0fc2a35 commit 2ac04d1
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,13 @@ As a community contributor, you will need to have an Atlas cluster to test Mongo
1. Create `secrets/credentials.json` file
1. Insert below json to the file with your configuration
```
{
"database": "database_name",
"user": "user",
"password": "password",
"cluster_url": "cluster_url"
{
"database": "database_name",
"user": "username",
"password": "password",
"connection_string": "mongodb+srv://cluster0.abcd1.mongodb.net/",
"replica_set": "atlas-abcdefg-shard-0",
"auth_source": "auth_database"
}
```
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ plugins {
id 'airbyte-docker'
id 'airbyte-integration-test-java'
id 'airbyte-connector-acceptance-test'
id 'org.jetbrains.kotlin.jvm' version '1.9.0'
}

application {
Expand All @@ -20,7 +21,35 @@ dependencies {

implementation 'org.mongodb:mongodb-driver-sync:4.10.2'

testImplementation "org.jetbrains.kotlinx:kotlinx-cli:0.3.5"

integrationTestJavaImplementation project(':airbyte-integrations:bases:standard-source-test')
integrationTestJavaImplementation project(':airbyte-integrations:connectors:source-mongodb-internal-poc')
integrationTestJavaImplementation files(project(':airbyte-integrations:bases:base-java').airbyteDocker.outputs)
}

/*
* Executes the script that generates test data and inserts it into the provided database/collection.
*
* To execute this task, use the following command:
*
* ./gradlew :airbyte-integrations:connectors:source-mongodb-internal-poc:generateTestData -PconnectionString=<connection string> -PreplicaSet=<replica set> -PdatabaseName=<database name> -PcollectionName=<collection name> -Pusername=<username>
*
* Optionally, you can provide -PnumberOfDocuments to change the number of generated documents from the default (10,000).
*/
tasks.register('generateTestData', JavaExec) {
def arguments = ['--connection-string', connectionString,
'--database-name', databaseName,
'--collection-name', collectionName,
'--replica-set', replicaSet,
'--username', username]

if (project.hasProperty('numberOfDocuments')) {
arguments.addAll(['--number', numberOfDocuments])
}

classpath = sourceSets.test.runtimeClasspath
main 'io.airbyte.integrations.source.mongodb.internal.MongoDbInsertClient'
standardInput = System.in
args arguments
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
package io.airbyte.integrations.source.mongodb.internal

import io.airbyte.commons.json.Jsons
import kotlinx.cli.ArgParser
import kotlinx.cli.ArgType
import kotlinx.cli.default
import kotlinx.cli.required
import org.bson.BsonTimestamp
import org.bson.Document
import java.lang.System.currentTimeMillis

object MongoDbInsertClient {

@JvmStatic
fun main(args: Array<String>) {
val parser = ArgParser("MongoDb Insert Client")
val connectionString by parser.option(ArgType.String, fullName = "connection-string", shortName = "cs", description = "MongoDb Connection String").required()
val databaseName by parser.option(ArgType.String, fullName = "database-name", shortName = "d", description = "Database Name").required()
val collectionName by parser.option(ArgType.String, fullName = "collection-name", shortName = "cn", description = "Collection Name").required()
val replicaSet by parser.option(ArgType.String, fullName = "replica-set", shortName = "r", description = "Replica Set").required()
val username by parser.option(ArgType.String, fullName = "username", shortName = "u", description = "Username").required()
val numberOfDocuments by parser.option(ArgType.Int, fullName = "number", shortName = "n", description = "Number of documents to generate").default(10000)

parser.parse(args)

println("Enter password: ")
val password = readln()

var config = mapOf(MongoConstants.DATABASE_CONFIGURATION_KEY to databaseName,
MongoConstants.CONNECTION_STRING_CONFIGURATION_KEY to connectionString,
MongoConstants.AUTH_SOURCE_CONFIGURATION_KEY to "admin",
MongoConstants.REPLICA_SET_CONFIGURATION_KEY to replicaSet,
MongoConstants.USER_CONFIGURATION_KEY to username,
MongoConstants.PASSWORD_CONFIGURATION_KEY to password)

MongoConnectionUtils.createMongoClient(Jsons.deserialize(Jsons.serialize(config))).use { mongoClient ->
val documents = mutableListOf<Document>()
for (i in 0..numberOfDocuments) {
documents += Document().append("name", "Document $i")
.append("description", "This is document #$i")
.append("doubleField", i.toDouble())
.append("intField", i)
.append("objectField", mapOf("key" to "value"))
.append("timestamp", BsonTimestamp(currentTimeMillis()))
}

mongoClient.getDatabase(databaseName).getCollection(collectionName).insertMany(documents)
}

println("Inserted $numberOfDocuments document(s) to $databaseName.$collectionName")
}
}

0 comments on commit 2ac04d1

Please sign in to comment.