Fixes for the timeout on ES

v2/elasticsearch-common/src/main/java/com/google/cloud/teleport/v2/elasticsearch/options/ElasticsearchWriteOptions.java

@@ -64,6 +64,12 @@ public interface ElasticsearchWriteOptions extends PipelineOptions {
 
   void setBatchSizeBytes(Long batchSizeBytes);
 
+  @Description("Elastic search socket timeout. Default: 30 minutes")
+  @Default.Integer(1800000)
+  Integer getMaxSocketTimeout();
+
+  void setMaxSocketTimeout(Integer maxSocketTimeout);
+
   @Description("Optional: Max retry attempts, must be > 0, ex: 3. Default: no retries")
   Integer getMaxRetryAttempts();
 

v2/elasticsearch-common/src/main/java/com/google/cloud/teleport/v2/elasticsearch/transforms/WriteToElasticsearch.java

@@ -114,6 +114,7 @@ public PDone expand(PCollection<String> jsonStrings) {
             options().getIndex(),
             DOCUMENT_TYPE);
 
+    config = config.withSocketTimeout(options().getMaxSocketTimeout());
     // If username and password are not blank, use them instead of ApiKey
     if (StringUtils.isNotBlank(options().getElasticsearchUsername())
         && StringUtils.isNotBlank(options().getElasticsearchPassword())) {

v2/elasticsearch-common/src/main/java/com/google/cloud/teleport/v2/elasticsearch/utils/ElasticsearchIO.java

@@ -96,6 +96,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.cloud.ServiceOptions;
+import com.google.cloud.secretmanager.v1.AccessSecretVersionResponse;
+import com.google.cloud.secretmanager.v1.SecretManagerServiceClient;
+import com.google.cloud.secretmanager.v1.SecretVersionName;
+
+
 /**
  * Transforms for reading and writing data from/to Elasticsearch.
  *
@@ -469,8 +475,20 @@ RestClient createClient() throws IOException {
                 httpAsyncClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
       }
       if (getApiKey() != null) {
+        SecretManagerServiceClient client = SecretManagerServiceClient.create();
+        String projectId = ServiceOptions.getDefaultProjectId();
+        String secretId = getApiKey();
+        String versionId = "latest";
+        SecretVersionName secretVersionName = SecretVersionName.of(projectId, secretId, versionId);
+
+        AccessSecretVersionResponse response = client.accessSecretVersion(secretVersionName);
+        String plainTextSecret = response.getPayload().getData().toStringUtf8();
+        URL url = new URL(plainTextSecret);
+        String userInfo = url.getUserInfo();
+
+        String[] userInfoArray = userInfo.split(":");        
         restClientBuilder.setDefaultHeaders(
-            new Header[] {new BasicHeader("Authorization", "ApiKey " + getApiKey())});
+            new Header[] {new BasicHeader("Authorization", "ApiKey " + userInfoArray[1])});
       }
       if (getBearerToken() != null) {
         restClientBuilder.setDefaultHeaders(

v2/googlecloud-to-elasticsearch/src/main/resources/googlecloud-to-elasticsearch-command-spec.json

@@ -0,0 +1,7 @@
+{
+  "mainClass": "com.google.cloud.teleport.v2.elasticsearch.templates.BigQueryToElasticsearch",
+  "classPath": "/template/googlecloud-to-elasticsearch/*:/template/googlecloud-to-elasticsearch/libs/*:/template/googlecloud-to-elasticsearch/classes",
+  "defaultParameterValues": {
+    "labels": "{\"goog-dataflow-provided-template-type\":\"flex\", \"goog-dataflow-provided-template-name\":\"bigquery_to_elasticsearch\"}"
+  }
+}

v2/pom.xml

@@ -173,12 +173,12 @@
             <version>${mockito-core.version}</version>
             <scope>test</scope>
         </dependency>
-        <dependency>
+        <!-- <dependency>
             <groupId>com.google.cloud.teleport</groupId>
             <artifactId>integration-testing-lib</artifactId>
             <version>1.0-SNAPSHOT</version>
             <scope>test</scope>
-        </dependency>
+        </dependency> -->
 
 
         <!-- Misc -->