Repository: googlecloudplatform/dataflowtemplates
Issue: [Bug]: MongoDBtoBigQuery with UDF: ScriptObjectMirror cannot be cast to bson.Document

### Related Template(s)

MongoDB to BigQuery

### What happened?

When creating a job with MongoDB to BigQuery template and specify a UDF javascript file and function, the job fails to start. See relevant logs below.

When UDF file and function are not specified, jobs with the same template can run normally.

### Beam Version

Newer than 2.43.0

### Relevant log output

```shell
com.google.cloud.teleport.v2.common.UncaughtExceptionLogger - The template launch failed.
java.lang.ClassCastException: class org.openjdk.nashorn.api.scripting.ScriptObjectMirror cannot be cast to class org.bson.Document (org.openjdk.nashorn.api.scripting.ScriptObjectMirror and org.bson.Document are in unnamed module of loader 'app')
	at com.google.cloud.teleport.v2.mongodb.templates.MongoDbUtils.getTableFieldSchemaForUDF(MongoDbUtils.java:167)
	at com.google.cloud.teleport.v2.mongodb.templates.MongoDbToBigQuery.run(MongoDbToBigQuery.java:95)
	at com.google.cloud.teleport.v2.mongodb.templates.MongoDbToBigQuery.main(MongoDbToBigQuery.java:82)
```


Changed files:
v2/mongodb-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbUtils.java
v2/mongodb-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbToBigQueryIT.java
v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/JavascriptDocumentTransformer.java

diff --git a/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/JavascriptDocumentTransformer.java b/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/JavascriptDocumentTransformer.java
index a95016ed..85460f1a 100644
--- a/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/JavascriptDocumentTransformer.java
+++ b/v2/common/src/main/java/com/google/cloud/teleport/v2/transforms/JavascriptDocumentTransformer.java
@@ -147,15 +147,17 @@ public abstract class JavascriptDocumentTransformer {
         throw new RuntimeException("No udf was loaded");
       }
 
-      Object result = getInvocable().invokeFunction(functionName(), data);
+      Object result = getInvocable().invokeFunction(functionName(), data.toJson());
       if (result == null || ScriptObjectMirror.isUndefined(result)) {
         return null;
       } else if (result instanceof Document) {
         return (Document) result;
+      } else if (result instanceof String) {
+        return Document.parse(result.toString());
       } else {
         String className = result.getClass().getName();
         throw new RuntimeException(
-            "UDF Function did not return a String. Instead got: " + className);
+            "UDF Function did not return a valid Mongo Document. Instead got: " + className);
       }
     }
 
diff --git a/v2/mongodb-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbUtils.java b/v2/mongodb-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbUtils.java
index 0b7258e0..9e242f87 100644
--- a/v2/mongodb-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbUtils.java
+++ b/v2/mongodb-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbUtils.java
@@ -50,6 +50,7 @@ import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
 import org.apache.beam.sdk.io.fs.MatchResult.Status;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.CharStreams;
 import org.bson.Document;
+import org.openjdk.nashorn.api.scripting.ScriptObjectMirror;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -163,8 +164,23 @@ public class MongoDbUtils implements Serializable {
       throw new RuntimeException("No udf was loaded");
     }
 
-    Object result = invocable.invokeFunction(udfFunctionName, document);
-    Document doc = (Document) result;
+    Document doc;
+    Object result = invocable.invokeFunction(udfFunctionName, document.toJson());
+    if (result == null || ScriptObjectMirror.isUndefined(result)) {
+      return null;
+    } else if (result instanceof Document) {
+      doc = (Document) result;
+    } else if (result instanceof String) {
+      doc = Document.parse(result.toString());
+    } else {
+      String className = result.getClass().getName();
+      throw new RuntimeException(
+          "UDF Function did not return a valid Mongo Document. Instead got: "
+              + className
+              + ": "
+              + result);
+    }
+
     if (userOption.equals("FLATTEN")) {
       doc.forEach(
           (key, value) -> {
diff --git a/v2/mongodb-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbToBigQueryIT.java b/v2/mongodb-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbToBigQueryIT.java
index 6def5429..03ccb5a3 100644
--- a/v2/mongodb-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbToBigQueryIT.java
+++ b/v2/mongodb-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/mongodb/templates/MongoDbToBigQueryIT.java
@@ -132,6 +132,14 @@ public final class MongoDbToBigQueryIT extends TemplateTestBase {
     mongoDbClient.insertDocuments(collectionName, mongoDocuments);
 
     String bqTable = testName;
+    String udfFileName = "transform.js";
+    artifactClient.createArtifact(
+        "input/" + udfFileName,
+        "function transform(inJson) {\\n"
+            + "    var outJson = JSON.parse(inJson);\\n"
+            + "    outJson.udf = \\"out\\";\\n"
+            + "    return JSON.stringify(outJson);\\n"
+            + "}");
     List<Field> bqSchemaFields = new ArrayList<>();
     bqSchemaFields.add(Field.of("timestamp", StandardSQLTypeName.TIMESTAMP));
     mongoDocuments
@@ -148,7 +156,9 @@ public final class MongoDbToBigQueryIT extends TemplateTestBase {
             .addParameter(MONGO_DB, mongoDbClient.getDatabaseName())
             .addParameter(MONGO_COLLECTION, collectionName)
             .addParameter(BIGQUERY_TABLE, toTableSpec(table))
-            .addParameter(USER_OPTION, "FLATTEN");
+            .addParameter(USER_OPTION, "FLATTEN")
+            .addParameter("javascriptDocumentTransformGcsPath", getGcsPath("input/" + udfFileName))
+            .addParameter("javascriptDocumentTransformFunctionName", "transform");
 
     // Act
     LaunchInfo info = launchTemplate(options);
@@ -168,6 +178,7 @@ public final class MongoDbToBigQueryIT extends TemplateTestBase {
         mongoDocument -> {
           JSONObject mongoDbJson = new JSONObject(mongoDocument.toJson());
           String mongoId = mongoDbJson.getJSONObject(MONGO_DB_ID).getString("$oid");
+          mongoDbJson.put("udf", "out");
           mongoDbJson.put(MONGO_DB_ID, mongoId);
           mongoMap.put(mongoId, mongoDbJson);
         });
@@ -212,6 +223,7 @@ public final class MongoDbToBigQueryIT extends TemplateTestBase {
       }
       mongoDocumentKeys.add(randomFieldName.toLowerCase());
     }
+    mongoDocumentKeys.add("udf");
 
     for (int i = 0; i < numDocuments; i++) {
       Document randomDocument = new Document().append(MONGO_DB_ID, new ObjectId());
@@ -220,6 +232,7 @@ public final class MongoDbToBigQueryIT extends TemplateTestBase {
         randomDocument.append(
             mongoDocumentKeys.get(j), RandomStringUtils.randomAlphanumeric(0, 20));
       }
+      randomDocument.append("udf", "in");
 
       mongoDocuments.add(randomDocument);
     }