feat: add MongoDB driver handshake metadata to MongoDBResourceManager

Signed-off-by: Alex Bevilacqua <alex@alexbevi.com>

Author: alexbevi

it/mongodb/src/main/java/org/apache/beam/it/mongodb/MongoDBResourceManager.java

@@ -20,6 +20,8 @@
 import static org.apache.beam.it.mongodb.MongoDBResourceManagerUtils.checkValidCollectionName;
 import static org.apache.beam.it.mongodb.MongoDBResourceManagerUtils.generateDatabaseName;
 
+import com.mongodb.ConnectionString;
+import com.mongodb.MongoDriverInformation;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
@@ -57,6 +59,10 @@ public class MongoDBResourceManager extends TestContainerResourceManager<MongoDB
 
   private static final String DEFAULT_MONGODB_CONTAINER_NAME = "mongo";
 
+  @VisibleForTesting
+  static final MongoDriverInformation DRIVER_INFO =
+      MongoDriverInformation.builder().driverName("Apache Beam").build();
+
   // A list of available MongoDB Docker image tags can be found at
   // https://hub.docker.com/_/mongo/tags
   private static final String DEFAULT_MONGODB_CONTAINER_TAG = "4.0.18";
@@ -88,7 +94,10 @@ private MongoDBResourceManager(Builder builder) {
         usingStaticDatabase ? builder.databaseName : generateDatabaseName(builder.testId);
     this.connectionString =
         String.format("mongodb://%s:%d", this.getHost(), this.getPort(MONGODB_INTERNAL_PORT));
-    this.mongoClient = mongoClient == null ? MongoClients.create(connectionString) : mongoClient;
+    this.mongoClient =
+        mongoClient == null
+            ? MongoClients.create(new ConnectionString(connectionString), DRIVER_INFO)
+            : mongoClient;
   }
 
   public static Builder builder(String testId) {

it/mongodb/src/test/java/org/apache/beam/it/mongodb/MongoDBResourceManagerTest.java

@@ -71,6 +71,11 @@ public void setUp() {
         new MongoDBResourceManager(mongoClient, container, MongoDBResourceManager.builder(TEST_ID));
   }
 
+  @Test
+  public void testDriverInfoHasExpectedName() {
+    assertThat(MongoDBResourceManager.DRIVER_INFO.getDriverNames()).contains("Apache Beam");
+  }
+
   @Test
   public void testCreateResourceManagerBuilderReturnsMongoDBResourceManager() {
     assertThat(

sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java

@@ -23,6 +23,7 @@
 import com.google.auto.value.AutoValue;
 import com.mongodb.ConnectionString;
 import com.mongodb.MongoClientSettings;
+import com.mongodb.MongoDriverInformation;
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
 import com.mongodb.client.MongoCursor;
@@ -119,6 +120,9 @@
  */
 public class MongoDbGridFSIO {
 
+  private static final MongoDriverInformation DRIVER_INFO =
+      MongoDriverInformation.builder().driverName("Apache Beam").build();
+
   /** Callback for the parser to use to submit data. */
   public interface ParserCallback<T> extends Serializable {
     /** Output the object. The default timestamp will be the GridFSFile creation timestamp. */
@@ -203,13 +207,13 @@ static ConnectionConfiguration create(
 
     MongoClient setupMongo() {
       if (uri() == null) {
-        return MongoClients.create();
+        return MongoClients.create(MongoClientSettings.builder().build(), DRIVER_INFO);
       }
       MongoClientSettings settings =
           MongoClientSettings.builder()
               .applyConnectionString(new ConnectionString(Preconditions.checkStateNotNull(uri())))
               .build();
-      return MongoClients.create(settings);
+      return MongoClients.create(settings, DRIVER_INFO);
     }
 
     GridFSBucket setupGridFS(MongoClient mongo) {

sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java

@@ -27,6 +27,7 @@
 import com.mongodb.MongoClientSettings;
 import com.mongodb.MongoClientSettings.Builder;
 import com.mongodb.MongoCommandException;
+import com.mongodb.MongoDriverInformation;
 import com.mongodb.client.AggregateIterable;
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
@@ -145,6 +146,9 @@ public class MongoDbIO {
 
   private static final Logger LOG = LoggerFactory.getLogger(MongoDbIO.class);
 
+  private static final MongoDriverInformation DRIVER_INFO =
+      MongoDriverInformation.builder().driverName("Apache Beam").build();
+
   public static final String ERROR_MSG_QUERY_FN =
       " class is not supported. "
           + "Please provide one of the predefined classes in the MongoDbIO package "
@@ -427,7 +431,7 @@ long getDocumentCount() {
                   spec.ignoreSSLCertificate())
               .applyConnectionString(new ConnectionString(uri))
               .build();
-      try (MongoClient mongoClient = MongoClients.create(settings)) {
+      try (MongoClient mongoClient = MongoClients.create(settings, DRIVER_INFO)) {
         return getDocumentCount(mongoClient, database, collection);
       } catch (Exception e) {
         return -1;
@@ -459,7 +463,7 @@ public long getEstimatedSizeBytes(PipelineOptions pipelineOptions) {
                   spec.ignoreSSLCertificate())
               .applyConnectionString(new ConnectionString(uri))
               .build();
-      try (MongoClient mongoClient = MongoClients.create(settings)) {
+      try (MongoClient mongoClient = MongoClients.create(settings, DRIVER_INFO)) {
         try {
           return getEstimatedSizeBytes(mongoClient, database, collection);
         } catch (MongoCommandException exception) {
@@ -496,7 +500,7 @@ public List<BoundedSource<Document>> split(
                   spec.ignoreSSLCertificate())
               .applyConnectionString(new ConnectionString(uri))
               .build();
-      try (MongoClient mongoClient = MongoClients.create(settings)) {
+      try (MongoClient mongoClient = MongoClients.create(settings, DRIVER_INFO)) {
         MongoDatabase mongoDatabase = mongoClient.getDatabase(database);
 
         List<Document> splitKeys;
@@ -812,7 +816,7 @@ private MongoClient createClient(Read spec) {
                   spec.ignoreSSLCertificate())
               .applyConnectionString(new ConnectionString(uri))
               .build();
-      return MongoClients.create(settings);
+      return MongoClients.create(settings, DRIVER_INFO);
     }
   }
 
@@ -1012,7 +1016,7 @@ public void createMongoClient() {
                     spec.ignoreSSLCertificate())
                 .applyConnectionString(new ConnectionString(uri))
                 .build();
-        client = MongoClients.create(settings);
+        client = MongoClients.create(settings, DRIVER_INFO);
       }
 
       @StartBundle