✨ Source MongoDB Internal POC: Generate Test Data

airbyte-integrations/connectors/source-mongodb-internal-poc/README.md

@@ -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"
        }
       ```
 

airbyte-integrations/connectors/source-mongodb-internal-poc/build.gradle

@@ -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 {
@@ -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
+}

airbyte-integrations/connectors/source-mongodb-internal-poc/src/test/kotlin/MongoDbInsertClient.kt

@@ -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: ")
+        var 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 ->
+            var 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")
+    }
+}