Merge pull request #89 from microsoft/feature/62-poll-for-completion

Consumer: poll the data destination for completion

Author: paullatzelsperger

.github/workflows/verify.yaml

@@ -16,6 +16,7 @@ jobs:
       - name: Gradle Test Core
         env:
           AZ_STORAGE_SAS: ${{ secrets.AZ_STORAGE_SAS }}
+          AZ_STORAGE_KEY: ${{ secrets.AZ_STORAGE_KEY }}
           NIFI_API_AUTH: ${{ secrets.NIFI_API_AUTH }}
           S3_SECRET_ACCESS_KEY: ${{ secrets.S3_SECRET_ACCESS_KEY }}
           S3_ACCESS_KEY_ID: ${{ secrets.S3_ACCESS_KEY_ID }}

common/build.gradle.kts

@@ -5,10 +5,19 @@
 
 plugins {
     `java-library`
+    `java-test-fixtures`
 }
 
+val storageBlobVersion: String by project;
+
 dependencies {
     api(project(":spi"))
-    
+
+    testFixturesApi(platform("com.amazonaws:aws-java-sdk-bom:1.11.1018"))
+    testFixturesApi("com.amazonaws:aws-java-sdk-s3")
+    testFixturesApi("com.azure:azure-storage-blob:${storageBlobVersion}")
+
+    testFixturesImplementation("org.junit.jupiter:junit-jupiter-api:5.5.2")
+    testFixturesRuntimeOnly("org.junit.jupiter:junit-jupiter-engine:5.5.2")
 }
 

common/src/testFixtures/java/com/microsoft/dagx/common/testfixtures/AbstractAzureBlobTest.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ *  All rights reserved.
+ *
+ */
+
+package com.microsoft.dagx.common.testfixtures;
+
+import com.azure.core.credential.AzureSasCredential;
+import com.azure.storage.blob.BlobContainerClient;
+import com.azure.storage.blob.BlobServiceClient;
+import com.azure.storage.blob.BlobServiceClientBuilder;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+
+import java.io.File;
+import java.util.Objects;
+import java.util.UUID;
+
+import static com.microsoft.dagx.common.ConfigurationFunctions.propOrEnv;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class AbstractAzureBlobTest {
+
+    protected static final String accountName = "dagxblobstoreitest";
+    protected static BlobServiceClient blobServiceClient;
+    protected String containerName;
+    protected boolean reuseClient = true;
+    protected String testRunId;
+
+    @BeforeEach
+    public void setupClient() {
+
+        testRunId = UUID.randomUUID().toString();
+        containerName = "fetch-azure-processor-" + testRunId;
+
+        if (blobServiceClient == null || !reuseClient) {
+            final String accountSas = getSasToken();
+            blobServiceClient = new BlobServiceClientBuilder().credential(new AzureSasCredential(accountSas)).endpoint("https://" + accountName + ".blob.core.windows.net").buildClient();
+        }
+
+        if (blobServiceClient.getBlobContainerClient(containerName).exists()) {
+            fail("Container " + containerName + " already exists - tests  will fail!");
+        }
+
+        //create container
+        BlobContainerClient blobContainerClient = blobServiceClient.createBlobContainer(containerName);
+        if (!blobContainerClient.exists()) {
+            fail("Setup incomplete, tests will fail");
+
+        }
+    }
+
+    @AfterEach
+    public void teardown() {
+        try {
+            blobServiceClient.deleteBlobContainer(containerName);
+        } catch (Exception ex) {
+            fail("teardown failed, subsequent tests might fail as well!");
+        }
+    }
+
+    @NotNull
+    protected String getSasToken() {
+        return Objects.requireNonNull(propOrEnv("AZ_STORAGE_SAS", null), "AZ_STORAGE_SAS");
+    }
+
+    protected void putBlob(String name, File file) {
+        blobServiceClient.getBlobContainerClient(containerName)
+                .getBlobClient(name)
+                .uploadFromFile(file.getAbsolutePath(), true);
+    }
+}

external/nifi/processors/src/test/java/microsoft/dagx/transfer/nifi/processors/AbstractS3Test.java renamed to common/src/testFixtures/java/com/microsoft/dagx/common/testfixtures/AbstractS3Test.java

@@ -4,7 +4,7 @@
  *
  */
 
-package microsoft.dagx.transfer.nifi.processors;
+package com.microsoft.dagx.common.testfixtures;
 
 import com.amazonaws.auth.AWSCredentials;
 import com.amazonaws.auth.AWSStaticCredentialsProvider;
@@ -22,9 +22,12 @@
 import static com.microsoft.dagx.common.ConfigurationFunctions.propOrEnv;
 import static org.junit.jupiter.api.Assertions.fail;
 
-public class AbstractS3Test {
+/**
+ * Base class for tests that need an S3 bucket created and deleted on every test run.
+ */
+public abstract class AbstractS3Test {
 
-    protected final static String REGION = System.getProperty("it.aws.region", Regions.US_EAST_1.getName());
+    protected final static String region = System.getProperty("it.aws.region", Regions.US_EAST_1.getName());
     // Adding REGION to bucket prevents errors of
     //      "A conflicting conditional operation is currently in progress against this resource."
     // when bucket is rapidly added/deleted and consistency propagation causes this error.
@@ -36,12 +39,17 @@ public class AbstractS3Test {
 
     @BeforeEach
     public void setupClient() {
-        bucketName = "test-bucket-" + System.currentTimeMillis() + "-" + REGION;
+        bucketName = createBucketName();
         credentials = getCredentials();
-        client = AmazonS3ClientBuilder.standard().withCredentials(new AWSStaticCredentialsProvider(credentials)).withRegion(REGION).build();
+        client = AmazonS3ClientBuilder.standard().withCredentials(new AWSStaticCredentialsProvider(credentials)).withRegion(region).build();
         createBucket(bucketName);
     }
 
+    @NotNull
+    protected String createBucketName() {
+        return "test-bucket-" + System.currentTimeMillis() + "-" + region;
+    }
+
     @AfterEach
     void cleanup() {
         deleteBucket(bucketName);
@@ -59,9 +67,9 @@ protected void createBucket(String bucketName) {
             fail("Bucket " + bucketName + " exists. Choose a different bucket name to continue test");
         }
 
-        CreateBucketRequest request = AbstractS3Test.REGION.contains("east")
+        CreateBucketRequest request = AbstractS3Test.region.contains("east")
                 ? new CreateBucketRequest(bucketName) // See https://github.com/boto/boto3/issues/125
-                : new CreateBucketRequest(bucketName, AbstractS3Test.REGION);
+                : new CreateBucketRequest(bucketName, AbstractS3Test.region);
         client.createBucket(request);
 
         if (!client.doesBucketExistV2(bucketName)) {
@@ -84,7 +92,7 @@ protected void deleteBucket(String bucketName) {
                 }
 
                 if (objectListing.isTruncated()) {
-                    objectListing = client.listNextBatchOfObjects(objectListing);
+                    objectListing = client.listNextBatchOfObjects(objectListing);/**/
                 } else {
                     break;
                 }

external/nifi/processors/src/test/java/microsoft/dagx/transfer/nifi/processors/TestUtils.java renamed to common/src/testFixtures/java/com/microsoft/dagx/common/testfixtures/TestUtils.java

@@ -4,7 +4,7 @@
  *
  */
 
-package microsoft.dagx.transfer.nifi.processors;
+package com.microsoft.dagx.common.testfixtures;
 
 import java.io.File;
 import java.net.URI;
@@ -15,7 +15,7 @@
 import static org.junit.jupiter.api.Assertions.fail;
 
 public class TestUtils {
-    protected final static String SAMPLE_FILE_RESOURCE_NAME = "hello.txt";
+    public final static String SAMPLE_FILE_RESOURCE_NAME = "hello.txt";
 
     public static Path getResourcePath(String resourceName) {
         Path path = null;

extensions/transfer/transfer-core/build.gradle.kts

@@ -11,6 +11,7 @@ plugins {
 dependencies {
     api(project(":spi"))
 
+    testImplementation(project(":extensions:transfer:transfer-store-memory"))
 
 }
 

extensions/transfer/transfer-core/src/main/java/com/microsoft/dagx/transfer/core/CoreTransferExtension.java

@@ -18,13 +18,16 @@
 import com.microsoft.dagx.spi.transfer.store.TransferProcessStore;
 import com.microsoft.dagx.spi.types.TypeManager;
 import com.microsoft.dagx.spi.types.domain.transfer.DataRequest;
+import com.microsoft.dagx.spi.types.domain.transfer.StatusCheckerRegistry;
 import com.microsoft.dagx.transfer.core.flow.DataFlowManagerImpl;
 import com.microsoft.dagx.transfer.core.protocol.provider.RemoteMessageDispatcherRegistryImpl;
 import com.microsoft.dagx.transfer.core.provision.ProvisionManagerImpl;
 import com.microsoft.dagx.transfer.core.provision.ResourceManifestGeneratorImpl;
 import com.microsoft.dagx.transfer.core.transfer.ExponentialWaitStrategy;
 import com.microsoft.dagx.transfer.core.transfer.TransferProcessManagerImpl;
 
+import java.util.Set;
+
 /**
  * Provides core data transfer services to the system.
  */
@@ -44,7 +47,7 @@ public LoadPhase phase() {
 
     @Override
     public void initialize(ServiceExtensionContext context) {
-        this.monitor = context.getMonitor();
+        monitor = context.getMonitor();
         this.context = context;
 
         var typeManager = context.getTypeManager();
@@ -60,19 +63,23 @@ public void initialize(ServiceExtensionContext context) {
         var manifestGenerator = new ResourceManifestGeneratorImpl();
         context.registerService(ResourceManifestGenerator.class, manifestGenerator);
 
+        var statusCheckerRegistry = new StatusCheckerRegistryImpl();
+        context.registerService(StatusCheckerRegistry.class, statusCheckerRegistry);
+
         var vault = context.getService(Vault.class);
 
         provisionManager = new ProvisionManagerImpl(vault, typeManager, monitor);
         context.registerService(ProvisionManager.class, provisionManager);
 
-        var waitStrategy = context.hasService(TransferWaitStrategy.class) ? context.getService(TransferWaitStrategy.class) :new ExponentialWaitStrategy(DEFAULT_ITERATION_WAIT);
-        
+        var waitStrategy = context.hasService(TransferWaitStrategy.class) ? context.getService(TransferWaitStrategy.class) : new ExponentialWaitStrategy(DEFAULT_ITERATION_WAIT);
+
         processManager = TransferProcessManagerImpl.Builder.newInstance()
                 .waitStrategy(waitStrategy)
                 .manifestGenerator(manifestGenerator)
                 .dataFlowManager(dataFlowManager)
                 .provisionManager(provisionManager)
                 .dispatcherRegistry(dispatcherRegistry)
+                .statusCheckerRegistry(statusCheckerRegistry)
                 .monitor(monitor)
                 .build();
 
@@ -81,6 +88,11 @@ public void initialize(ServiceExtensionContext context) {
         monitor.info("Initialized Core Transfer extension");
     }
 
+    @Override
+    public Set<String> provides() {
+        return Set.of("dagx:statuschecker", "dagx:dispatcher", "dagx:manifestgenerator");
+    }
+
     @Override
     public void start() {
         var transferProcessStore = context.getService(TransferProcessStore.class);

extensions/transfer/transfer-core/src/main/java/com/microsoft/dagx/transfer/core/StatusCheckerRegistryImpl.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ *  All rights reserved.
+ *
+ */
+
+package com.microsoft.dagx.transfer.core;
+
+import com.microsoft.dagx.spi.types.domain.transfer.ProvisionedDataDestinationResource;
+import com.microsoft.dagx.spi.types.domain.transfer.StatusChecker;
+import com.microsoft.dagx.spi.types.domain.transfer.StatusCheckerRegistry;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class StatusCheckerRegistryImpl implements StatusCheckerRegistry {
+    private final Map<Class<? extends ProvisionedDataDestinationResource>, StatusChecker<? extends ProvisionedDataDestinationResource>> inMemoryMap;
+
+    public StatusCheckerRegistryImpl() {
+        inMemoryMap = new HashMap<>();
+    }
+
+    @Override
+    public void register(Class<? extends ProvisionedDataDestinationResource> provisionedResourceClass, StatusChecker<? extends ProvisionedDataDestinationResource> statusChecker) {
+        inMemoryMap.put(provisionedResourceClass, statusChecker);
+    }
+
+    @Override
+    public <T extends ProvisionedDataDestinationResource> StatusChecker<T> resolve(Class<? extends ProvisionedDataDestinationResource> provisionedResourceClass) {
+        return (StatusChecker<T>) inMemoryMap.get(provisionedResourceClass);
+    }
+
+    @Override
+    public <T extends ProvisionedDataDestinationResource> StatusChecker<T> resolve(T resource) {
+        return resolve(resource.getClass());
+    }
+}

extensions/transfer/transfer-core/src/main/java/com/microsoft/dagx/transfer/core/provision/ProvisionManagerImpl.java

@@ -76,8 +76,9 @@ void onDestinationResource(ProvisionedDataDestinationResource destinationResourc
             return;
         }
 
-
-        transferProcess.getDataRequest().updateDestination(destinationResource.createDataDestination());
+        if (!destinationResource.isError()) {
+            transferProcess.getDataRequest().updateDestination(destinationResource.createDataDestination());
+        }
 
         if (secretToken != null) {
             String keyName = destinationResource.getResourceName();

extensions/transfer/transfer-core/src/main/java/com/microsoft/dagx/transfer/core/transfer/ExponentialWaitStrategy.java

@@ -11,17 +11,17 @@
  * Implements an exponential backoff strategy for failed iterations.
  */
 public class ExponentialWaitStrategy implements TransferWaitStrategy {
+    private final long successWaitPeriodMillis;
     private int errorCount = 0;
-    private long successWaitPeriod;
 
 
-    public ExponentialWaitStrategy(long successWaitPeriod) {
-        this.successWaitPeriod = successWaitPeriod;
+    public ExponentialWaitStrategy(long successWaitPeriodMillis) {
+        this.successWaitPeriodMillis = successWaitPeriodMillis;
     }
 
     @Override
     public long waitForMillis() {
-        return successWaitPeriod;
+        return successWaitPeriodMillis;
     }
 
     @Override
@@ -33,7 +33,7 @@ public void success() {
     public long retryInMillis() {
         errorCount++;
         double exponentialMultiplier = Math.pow(2.0, errorCount - 1);
-        double result = exponentialMultiplier * successWaitPeriod;
+        double result = exponentialMultiplier * successWaitPeriodMillis;
         return (long) Math.min(result, Long.MAX_VALUE);
     }
 }