From 5ff6419c9f06d41927a894133154090665b2026c Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Sat, 3 May 2025 11:41:06 -0400 Subject: [PATCH 1/5] Auto-download Astra DB SCB --- .gitignore | 2 + README.md | 1 + RELEASE.md | 3 + pom.xml | 6 + .../TargetUpsertRunDetailsStatement.java | 8 +- .../datastax/cdm/data/AstraDevOpsClient.java | 355 ++++++++++++++++++ .../cdm/properties/KnownProperties.java | 36 ++ .../datastax/cdm/job/ConnectionFetcher.scala | 38 +- src/resources/cdm-detailed.properties | 31 ++ .../cdm/data/AstraDevOpsClientTest.java | 201 ++++++++++ .../cdm/job/ConnectionFetcherTest.java | 88 ++++- 11 files changed, 757 insertions(+), 12 deletions(-) create mode 100644 src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java create mode 100644 src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java diff --git a/.gitignore b/.gitignore index dcbb3ad4..420cce8e 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,5 @@ SIT/**/output/* .settings/* src/main/main.iml src/test/test.iml + +**/.claude/settings.local.json diff --git a/README.md b/README.md index bfc0d29f..71875221 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,7 @@ spark-submit --properties-file cdm.properties \ - Fully containerized (Docker and K8s friendly) - SSL Support (including custom cipher algorithms) - Migrate from any Cassandra `Origin` ([Apache Cassandra®](https://cassandra.apache.org) / [DataStax Enterprise™](https://www.datastax.com/products/datastax-enterprise) / [DataStax Astra DB™](https://www.datastax.com/products/datastax-astra)) to any Cassandra `Target` ([Apache Cassandra®](https://cassandra.apache.org) / [DataStax Enterprise™](https://www.datastax.com/products/datastax-enterprise) / [DataStax Astra DB™](https://www.datastax.com/products/datastax-astra)) +- Automatic download of Secure Connect Bundles for Astra DB using the DevOps API - Supports migration/validation from and to [Azure Cosmos Cassandra](https://learn.microsoft.com/en-us/azure/cosmos-db/cassandra) - Validate migration accuracy and performance using a smaller randomized data-set - Supports adding custom fixed `writetime` and/or `ttl` diff --git a/RELEASE.md b/RELEASE.md index 13d8a5a7..7ba93f34 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,5 +1,8 @@ # Release Notes +## [5.3.0] - 2025-05-05 +- Auto-download Astra DB Secure Connect Bundle (SCB) when connecting to Astra DB. + ## [5.2.3] - 2025-04-15 - Randomized the pending token-range list returned by the `trackRun` feature (when rerunning a previously incomplete job) for better load distribution across the cluster. diff --git a/pom.xml b/pom.xml index c09d7de7..60db27a7 100644 --- a/pom.xml +++ b/pom.xml @@ -175,6 +175,12 @@ ${mockito.version} test + + org.mockito + mockito-junit-jupiter + ${mockito.version} + test + diff --git a/src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java b/src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java index 855785ee..04d08717 100644 --- a/src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java +++ b/src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java @@ -114,12 +114,8 @@ public Collection getPendingPartitions(long prevRunId, JobType j final List pendingParts = new ArrayList(); // Use an array of statuses for iteration - String[] statuses = { - TrackRun.RUN_STATUS.NOT_STARTED.toString(), - TrackRun.RUN_STATUS.STARTED.toString(), - TrackRun.RUN_STATUS.FAIL.toString(), - TrackRun.RUN_STATUS.DIFF.toString() - }; + String[] statuses = { TrackRun.RUN_STATUS.NOT_STARTED.toString(), TrackRun.RUN_STATUS.STARTED.toString(), + TrackRun.RUN_STATUS.FAIL.toString(), TrackRun.RUN_STATUS.DIFF.toString() }; for (String status : statuses) { pendingParts.addAll(getPartitionsByStatus(prevRunId, status, jobType)); } diff --git a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java new file mode 100644 index 00000000..2bb0e144 --- /dev/null +++ b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java @@ -0,0 +1,355 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.datastax.cdm.data; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.time.Duration; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.datastax.cdm.properties.IPropertyHelper; +import com.datastax.cdm.properties.KnownProperties; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +/** + * Client for interacting with the Astra DevOps API to download secure connect bundles. This client supports downloading + * different types of SCBs (default, regional, custom domain). + */ +public class AstraDevOpsClient { + private static final Logger logger = LoggerFactory.getLogger(AstraDevOpsClient.class.getName()); + private static final String ASTRA_API_BASE_URL = "https://api.astra.datastax.com"; + private static final String SCB_API_PATH = "/v2/databases/%s/secureBundleURL"; + private static final Duration HTTP_TIMEOUT = Duration.ofSeconds(30); + + private final IPropertyHelper propertyHelper; + private final HttpClient httpClient; + private final ObjectMapper objectMapper; + + /** + * Creates a new AstraDevOpsClient instance. + * + * @param propertyHelper + * The property helper for accessing configuration + */ + public AstraDevOpsClient(IPropertyHelper propertyHelper) { + this.propertyHelper = propertyHelper; + this.httpClient = HttpClient.newBuilder().connectTimeout(HTTP_TIMEOUT).build(); + this.objectMapper = new ObjectMapper(); + } + + /** + * Downloads a secure connect bundle for the specified side (ORIGIN or TARGET). + * + * @param side + * The side (ORIGIN or TARGET) + * + * @return The path to the downloaded SCB file, or null if fails + * + * @throws IOException + * If an error occurs during the download process + */ + public String downloadSecureBundle(PKFactory.Side side) throws IOException { + String token = getAstraToken(side); + String databaseId = getAstraDatabaseId(side); + String scbType = getScbType(side); + + if (token == null || token.isEmpty() || databaseId == null || databaseId.isEmpty()) { + logger.warn("Missing required Astra parameters for {} (token or database ID)", side); + return null; + } + + logger.info("Auto-downloading secure connect bundle for {} database ID: {}, type: {}", side, databaseId, + scbType); + + try { + String jsonResponse = fetchSecureBundleUrlInfo(token, databaseId, "region".equals(scbType)); + + if (jsonResponse == null) { + logger.error("Failed to fetch secure bundle URL info for {}", side); + return null; + } + + String downloadUrl = extractDownloadUrl(jsonResponse, scbType, side); + + if (downloadUrl == null) { + logger.error("Could not extract download URL for {} bundle type: {}", side, scbType); + return null; + } + + return downloadBundleFile(downloadUrl, side); + + } catch (Exception e) { + logger.error("Error downloading secure bundle for {}: {}", side, e.getMessage()); + throw new IOException("Failed to download secure bundle", e); + } + } + + /** + * Gets the Astra token for the specified side. + * + * @param side + * The side (ORIGIN or TARGET) + * + * @return The Astra token + */ + private String getAstraToken(PKFactory.Side side) { + String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_TOKEN + : KnownProperties.TARGET_ASTRA_TOKEN; + + return propertyHelper.getAsString(property); + } + + /** + * Gets the Astra database ID for the specified side. + * + * @param side + * The side (ORIGIN or TARGET) + * + * @return The Astra database ID + */ + private String getAstraDatabaseId(PKFactory.Side side) { + String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_DATABASE_ID + : KnownProperties.TARGET_ASTRA_DATABASE_ID; + + return propertyHelper.getAsString(property); + } + + /** + * Gets the SCB type for the specified side. + * + * @param side + * The side (ORIGIN or TARGET) + * + * @return The SCB type ("default", "region", or "custom") + */ + private String getScbType(PKFactory.Side side) { + String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_SCB_TYPE + : KnownProperties.TARGET_ASTRA_SCB_TYPE; + + return propertyHelper.getAsString(property); + } + + /** + * Gets the region for regional SCBs. + * + * @param side + * The side (ORIGIN or TARGET) + * + * @return The region name + */ + private String getRegion(PKFactory.Side side) { + String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_SCB_REGION + : KnownProperties.TARGET_ASTRA_SCB_REGION; + + return propertyHelper.getAsString(property); + } + + /** + * Gets the custom domain for custom SCBs. + * + * @param side + * The side (ORIGIN or TARGET) + * + * @return The custom domain + */ + private String getCustomDomain(PKFactory.Side side) { + String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_SCB_CUSTOM_DOMAIN + : KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN; + + return propertyHelper.getAsString(property); + } + + /** + * Fetches secure bundle URL information from the Astra DevOps API. + * + * @param token + * The Astra token + * @param databaseId + * The database ID + * @param fetchAll + * Whether to fetch all regional bundles + * + * @return The JSON response as a string + * + * @throws IOException + * If an error occurs during the API call + * @throws InterruptedException + * If the API call is interrupted + */ + private String fetchSecureBundleUrlInfo(String token, String databaseId, boolean fetchAll) + throws IOException, InterruptedException { + + String apiUrl = ASTRA_API_BASE_URL + String.format(SCB_API_PATH, databaseId); + + if (fetchAll) { + apiUrl += "?all=true"; + } + + HttpRequest request = HttpRequest.newBuilder().uri(URI.create(apiUrl)) + .header("Authorization", "Bearer " + token).header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.noBody()).timeout(HTTP_TIMEOUT).build(); + + HttpResponse response = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + + if (response.statusCode() != 200) { + logger.error("Failed to fetch secure bundle URL. Status code: {}, Response: {}", response.statusCode(), + response.body()); + return null; + } + + return response.body(); + } + + /** + * Extracts the appropriate download URL from the API response based on the SCB type. + * + * @param jsonResponse + * The JSON response from the API + * @param scbType + * The SCB type + * @param side + * The side (ORIGIN or TARGET) + * + * @return The download URL + * + * @throws IOException + * If an error occurs parsing the response + */ + private String extractDownloadUrl(String jsonResponse, String scbType, PKFactory.Side side) throws IOException { + JsonNode rootNode = objectMapper.readTree(jsonResponse); + + switch (scbType.toLowerCase()) { + case "default": + // Default bundle URL extraction + if (rootNode.has("downloadURL")) { + return rootNode.get("downloadURL").asText(); + } + break; + + case "region": + // Regional bundle URL extraction + String region = getRegion(side); + if (region == null || region.isEmpty()) { + logger.error("Region is required for SCB type 'region' but was not specified"); + return null; + } + + if (rootNode.has("downloadURLs") && rootNode.get("downloadURLs").isArray()) { + for (JsonNode regionNode : rootNode.get("downloadURLs")) { + if (regionNode.has("region") && region.equalsIgnoreCase(regionNode.get("region").asText()) + && regionNode.has("downloadURL")) { + + return regionNode.get("downloadURL").asText(); + } + } + logger.error("Could not find download URL for region: {}", region); + } + break; + + case "custom": + // Custom domain bundle URL extraction + String customDomain = getCustomDomain(side); + if (customDomain == null || customDomain.isEmpty()) { + logger.error("Custom domain is required for SCB type 'custom' but was not specified"); + return null; + } + + if (rootNode.has("customDomainBundles") && rootNode.get("customDomainBundles").isArray()) { + for (JsonNode customNode : rootNode.get("customDomainBundles")) { + if (customNode.has("sniDomain") + && customDomain.equalsIgnoreCase(customNode.get("sniDomain").asText()) + && customNode.has("downloadURL")) { + + return customNode.get("downloadURL").asText(); + } + } + logger.error("Could not find download URL for custom domain: {}", customDomain); + } + break; + + default: + logger.error("Unknown SCB type: {}", scbType); + break; + } + + return null; + } + + /** + * Downloads the secure bundle file from the specified URL. + * + * @param downloadUrl + * The URL to download from + * @param side + * The side (ORIGIN or TARGET) + * + * @return The path to the downloaded file + * + * @throws IOException + * If an error occurs during download + * @throws InterruptedException + * If the download is interrupted + */ + private String downloadBundleFile(String downloadUrl, PKFactory.Side side) + throws IOException, InterruptedException { + + logger.info("Downloading secure bundle from URL: {}", downloadUrl); + + HttpRequest downloadRequest = HttpRequest.newBuilder().uri(URI.create(downloadUrl)).GET() + .timeout(Duration.ofMinutes(2)).build(); + + Path tempDir = Files.createTempDirectory("cdm-scb-" + UUID.randomUUID()); + String fileName = side.toString().toLowerCase() + "-secure-bundle.zip"; + Path filePath = tempDir.resolve(fileName); + + HttpResponse downloadResponse = httpClient.send(downloadRequest, + HttpResponse.BodyHandlers.ofInputStream()); + + if (downloadResponse.statusCode() != 200) { + throw new IOException("Failed to download secure bundle. Status code: " + downloadResponse.statusCode()); + } + + try (InputStream in = downloadResponse.body(); FileOutputStream out = new FileOutputStream(filePath.toFile())) { + + byte[] buffer = new byte[8192]; + int bytesRead; + while ((bytesRead = in.read(buffer)) != -1) { + out.write(buffer, 0, bytesRead); + } + } + + // Ensure the file is deleted when the JVM exits + filePath.toFile().deleteOnExit(); + + logger.info("Secure bundle downloaded successfully to: {}", filePath); + return filePath.toString(); + } +} diff --git a/src/main/java/com/datastax/cdm/properties/KnownProperties.java b/src/main/java/com/datastax/cdm/properties/KnownProperties.java index 3bcedd41..9b5a26af 100644 --- a/src/main/java/com/datastax/cdm/properties/KnownProperties.java +++ b/src/main/java/com/datastax/cdm/properties/KnownProperties.java @@ -48,6 +48,23 @@ public enum PropertyType { public static final String CONNECT_TARGET_USERNAME = "spark.cdm.connect.target.username"; public static final String CONNECT_TARGET_PASSWORD = "spark.cdm.connect.target.password"; + // ========================================================================== + // Astra DevOps API Parameters + // ========================================================================== + public static final String ORIGIN_ASTRA_TOKEN = "spark.cdm.connect.origin.astra.token"; + public static final String ORIGIN_ASTRA_DATABASE_ID = "spark.cdm.connect.origin.astra.database.id"; + public static final String ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB = "spark.cdm.connect.origin.astra.auto.download.scb"; + public static final String ORIGIN_ASTRA_SCB_TYPE = "spark.cdm.connect.origin.astra.scb.type"; + public static final String ORIGIN_ASTRA_SCB_REGION = "spark.cdm.connect.origin.astra.scb.region"; + public static final String ORIGIN_ASTRA_SCB_CUSTOM_DOMAIN = "spark.cdm.connect.origin.astra.scb.custom.domain"; + + public static final String TARGET_ASTRA_TOKEN = "spark.cdm.connect.target.astra.token"; + public static final String TARGET_ASTRA_DATABASE_ID = "spark.cdm.connect.target.astra.database.id"; + public static final String TARGET_ASTRA_AUTO_DOWNLOAD_SCB = "spark.cdm.connect.target.astra.auto.download.scb"; + public static final String TARGET_ASTRA_SCB_TYPE = "spark.cdm.connect.target.astra.scb.type"; + public static final String TARGET_ASTRA_SCB_REGION = "spark.cdm.connect.target.astra.scb.region"; + public static final String TARGET_ASTRA_SCB_CUSTOM_DOMAIN = "spark.cdm.connect.target.astra.scb.custom.domain"; + static { types.put(CONNECT_ORIGIN_HOST, PropertyType.STRING); defaults.put(CONNECT_ORIGIN_HOST, "localhost"); @@ -68,6 +85,25 @@ public enum PropertyType { defaults.put(CONNECT_TARGET_USERNAME, "cassandra"); types.put(CONNECT_TARGET_PASSWORD, PropertyType.STRING); defaults.put(CONNECT_TARGET_PASSWORD, "cassandra"); + + // Astra DevOps API parameters + types.put(ORIGIN_ASTRA_TOKEN, PropertyType.STRING); + types.put(ORIGIN_ASTRA_DATABASE_ID, PropertyType.STRING); + types.put(ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB, PropertyType.BOOLEAN); + defaults.put(ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB, "false"); + types.put(ORIGIN_ASTRA_SCB_TYPE, PropertyType.STRING); + defaults.put(ORIGIN_ASTRA_SCB_TYPE, "default"); + types.put(ORIGIN_ASTRA_SCB_REGION, PropertyType.STRING); + types.put(ORIGIN_ASTRA_SCB_CUSTOM_DOMAIN, PropertyType.STRING); + + types.put(TARGET_ASTRA_TOKEN, PropertyType.STRING); + types.put(TARGET_ASTRA_DATABASE_ID, PropertyType.STRING); + types.put(TARGET_ASTRA_AUTO_DOWNLOAD_SCB, PropertyType.BOOLEAN); + defaults.put(TARGET_ASTRA_AUTO_DOWNLOAD_SCB, "false"); + types.put(TARGET_ASTRA_SCB_TYPE, PropertyType.STRING); + defaults.put(TARGET_ASTRA_SCB_TYPE, "default"); + types.put(TARGET_ASTRA_SCB_REGION, PropertyType.STRING); + types.put(TARGET_ASTRA_SCB_CUSTOM_DOMAIN, PropertyType.STRING); } // ========================================================================== diff --git a/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala b/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala index 9bae3e48..4ff3126d 100644 --- a/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala +++ b/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala @@ -20,13 +20,49 @@ import com.datastax.spark.connector.cql.CassandraConnector import org.apache.spark.SparkConf import org.slf4j.{Logger, LoggerFactory} import com.datastax.cdm.data.DataUtility.generateSCB +import com.datastax.cdm.data.{AstraDevOpsClient, PKFactory} import com.datastax.cdm.data.PKFactory.Side // TODO: CDM-31 - add localDC configuration support -class ConnectionFetcher(propertyHelper: IPropertyHelper) extends Serializable { +class ConnectionFetcher(propertyHelper: IPropertyHelper, testAstraClient: AstraDevOpsClient = null) extends Serializable { val logger: Logger = LoggerFactory.getLogger(this.getClass.getName) + + // Return injected client for testing, or create a new one + private def getAstraClient(): AstraDevOpsClient = { + if (testAstraClient != null) testAstraClient else new AstraDevOpsClient(propertyHelper) + } def getConnectionDetails(side: Side): ConnectionDetails = { + // Check if auto-download is enabled for this side + val autoDownloadEnabled = if (Side.ORIGIN.equals(side)) { + propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB) + } else { + propertyHelper.getBoolean(KnownProperties.TARGET_ASTRA_AUTO_DOWNLOAD_SCB) + } + + // If auto-download is enabled, try to download the secure bundle + if (autoDownloadEnabled) { + try { + val astraClient = getAstraClient() + val downloadedScbPath = astraClient.downloadSecureBundle(side) + + if (downloadedScbPath != null && !downloadedScbPath.isEmpty) { + logger.info(s"Successfully auto-downloaded secure bundle for $side: $downloadedScbPath") + + // Update the property helper with the downloaded SCB path + if (Side.ORIGIN.equals(side)) { + propertyHelper.setProperty(KnownProperties.CONNECT_ORIGIN_SCB, downloadedScbPath) + } else { + propertyHelper.setProperty(KnownProperties.CONNECT_TARGET_SCB, downloadedScbPath) + } + } + } catch { + case e: Exception => + logger.warn(s"Failed to auto-download secure bundle for $side: ${e.getMessage}") + logger.debug("Auto-download failure details", e) + } + } + if (Side.ORIGIN.equals(side)) { ConnectionDetails( propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_SCB), diff --git a/src/resources/cdm-detailed.properties b/src/resources/cdm-detailed.properties index af7fde31..0ac55c10 100644 --- a/src/resources/cdm-detailed.properties +++ b/src/resources/cdm-detailed.properties @@ -58,6 +58,37 @@ spark.cdm.connect.target.port 9042 spark.cdm.connect.target.username cassandra spark.cdm.connect.target.password cassandra +#=========================================================================================================== +# Astra DevOps API Parameters for auto-downloading Secure Connect Bundles +# These parameters enable automatic download of secure connect bundles from the Astra DevOps API +# This eliminates the need to manually download and specify the SCB path +# +# spark.cdm.connect.origin.astra +# spark.cdm.connect.target.astra +# .token : The Astra DevOps API token. Required if auto-download is enabled. +# .database.id : The Astra Database ID. Required if auto-download is enabled. +# .auto.download.scb : Default is false. Set to true to enable auto-download of the SCB. +# .scb.type : Default is 'default'. Can be 'default', 'region', or 'custom'. +# .scb.region : The region name for regional secure bundles. Required when scb.type is 'region'. +# .scb.custom.domain : The custom domain for custom secure bundles. Required when scb.type is 'custom'. +#----------------------------------------------------------------------------------------------------------- + +# Origin Astra DevOps API configuration +#spark.cdm.connect.origin.astra.token your-astra-devops-api-token +#spark.cdm.connect.origin.astra.database.id your-astra-database-id +#spark.cdm.connect.origin.astra.auto.download.scb false +#spark.cdm.connect.origin.astra.scb.type default +#spark.cdm.connect.origin.astra.scb.region us-east-1 +#spark.cdm.connect.origin.astra.scb.custom.domain your-custom-domain.example.com + +# Target Astra DevOps API configuration +#spark.cdm.connect.target.astra.token your-astra-devops-api-token +#spark.cdm.connect.target.astra.database.id your-astra-database-id +#spark.cdm.connect.target.astra.auto.download.scb false +#spark.cdm.connect.target.astra.scb.type default +#spark.cdm.connect.target.astra.scb.region us-east-1 +#spark.cdm.connect.target.astra.scb.custom.domain your-custom-domain.example.com + #=========================================================================================================== # Details about the Origin Schema # diff --git a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java new file mode 100644 index 00000000..b74b0202 --- /dev/null +++ b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java @@ -0,0 +1,201 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.datastax.cdm.data; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.net.http.HttpClient; +import java.net.http.HttpResponse; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import com.datastax.cdm.properties.IPropertyHelper; +import com.datastax.cdm.properties.KnownProperties; +import com.fasterxml.jackson.databind.ObjectMapper; + +@ExtendWith(MockitoExtension.class) +class AstraDevOpsClientTest { + + @Mock + private IPropertyHelper propertyHelper; + + @Mock + private HttpClient httpClient; + + @Mock + private HttpResponse httpResponse; + + private AstraDevOpsClient client; + + @BeforeEach + void setUp() throws Exception { + // Create the client with mocked property helper + client = new AstraDevOpsClient(propertyHelper); + + // Use reflection to replace httpClient with our mock + Field httpClientField = AstraDevOpsClient.class.getDeclaredField("httpClient"); + httpClientField.setAccessible(true); + httpClientField.set(client, httpClient); + } + + @Test + void testDownloadSecureBundleWithNullToken() throws IOException { + // Setup + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn(null); + + // Test + String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); + + // Verify + assertNull(result); + verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN); + } + + @Test + void testDownloadSecureBundleWithEmptyToken() throws IOException { + // Setup + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn(""); + + // Test + String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); + + // Verify + assertNull(result); + verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN); + } + + @Test + void testDownloadSecureBundleWithNullDatabaseId() throws IOException { + // Setup + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn(null); + + // Test + String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); + + // Verify + assertNull(result); + verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN); + verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID); + } + + @Test + void testGetAstraToken() throws Exception { + // Setup + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn("test-token"); + + // Use reflection to access the private method + Method getAstraTokenMethod = AstraDevOpsClient.class.getDeclaredMethod("getAstraToken", PKFactory.Side.class); + getAstraTokenMethod.setAccessible(true); + + // Test + String token = (String) getAstraTokenMethod.invoke(client, PKFactory.Side.ORIGIN); + + // Verify + assertEquals("test-token", token); + } + + @Test + void testGetAstraDatabaseId() throws Exception { + // Setup + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); + + // Use reflection to access the private method + Method getAstraDatabaseIdMethod = AstraDevOpsClient.class.getDeclaredMethod("getAstraDatabaseId", PKFactory.Side.class); + getAstraDatabaseIdMethod.setAccessible(true); + + // Test + String dbId = (String) getAstraDatabaseIdMethod.invoke(client, PKFactory.Side.TARGET); + + // Verify + assertEquals("test-db-id", dbId); + } + + @Test + void testGetScbType() throws Exception { + // Setup + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("region"); + + // Use reflection to access the private method + Method getScbTypeMethod = AstraDevOpsClient.class.getDeclaredMethod("getScbType", PKFactory.Side.class); + getScbTypeMethod.setAccessible(true); + + // Test + String scbType = (String) getScbTypeMethod.invoke(client, PKFactory.Side.ORIGIN); + + // Verify + assertEquals("region", scbType); + } + + @Test + void testGetRegion() throws Exception { + // Setup + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_REGION)).thenReturn("us-east-1"); + + // Use reflection to access the private method + Method getRegionMethod = AstraDevOpsClient.class.getDeclaredMethod("getRegion", PKFactory.Side.class); + getRegionMethod.setAccessible(true); + + // Test + String region = (String) getRegionMethod.invoke(client, PKFactory.Side.TARGET); + + // Verify + assertEquals("us-east-1", region); + } + + @Test + void testGetCustomDomain() throws Exception { + // Setup + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_CUSTOM_DOMAIN)).thenReturn("custom.domain.com"); + + // Use reflection to access the private method + Method getCustomDomainMethod = AstraDevOpsClient.class.getDeclaredMethod("getCustomDomain", PKFactory.Side.class); + getCustomDomainMethod.setAccessible(true); + + // Test + String customDomain = (String) getCustomDomainMethod.invoke(client, PKFactory.Side.ORIGIN); + + // Verify + assertEquals("custom.domain.com", customDomain); + } + + @Test + void testExtractDownloadUrlDefaultType() throws Exception { + // Setup + String jsonResponse = "{ \"downloadURL\": \"https://example.com/bundle.zip\" }"; + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "default", PKFactory.Side.ORIGIN); + + // Verify + assertEquals("https://example.com/bundle.zip", url); + } + + // More tests for extractDownloadUrl with different SCB types can be added here +} \ No newline at end of file diff --git a/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java b/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java index 8b7d0d5a..fbe0bced 100644 --- a/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java +++ b/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java @@ -16,19 +16,28 @@ package com.datastax.cdm.job; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.lang.reflect.Field; + import org.apache.spark.SparkConf; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; import com.datastax.cdm.cql.CommonMocks; +import com.datastax.cdm.data.AstraDevOpsClient; import com.datastax.cdm.data.PKFactory; import com.datastax.cdm.properties.IPropertyHelper; import com.datastax.cdm.properties.KnownProperties; +@MockitoSettings(strictness = Strictness.LENIENT) public class ConnectionFetcherTest extends CommonMocks { @Mock @@ -37,6 +46,9 @@ public class ConnectionFetcherTest extends CommonMocks { @Mock private SparkConf conf; + @Mock + private AstraDevOpsClient astraClient; + private ConnectionFetcher cf; @BeforeEach @@ -45,23 +57,89 @@ public void setup() { commonSetupWithoutDefaultClassVariables(); MockitoAnnotations.openMocks(this); - cf = new ConnectionFetcher(propertyHelper); + cf = new ConnectionFetcher(propertyHelper, null); + + // Default values for all tests + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_HOST)).thenReturn("origin_host"); + when(propertyHelper.getAsString(KnownProperties.CONNECT_TARGET_HOST)).thenReturn("target_host"); + + // Default - auto-download disabled for both ORIGIN and TARGET + when(propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(false); + when(propertyHelper.getBoolean(KnownProperties.TARGET_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(false); } @Test public void getConnectionDetailsOrigin() { - when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_HOST)).thenReturn("origin_host"); - when(propertyHelper.getAsString(KnownProperties.CONNECT_TARGET_HOST)).thenReturn("target_host"); ConnectionDetails cd = cf.getConnectionDetails(PKFactory.Side.ORIGIN); assertEquals("origin_host", cd.host()); } @Test public void getConnectionDetailsTarget() { - when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_HOST)).thenReturn("origin_host"); - when(propertyHelper.getAsString(KnownProperties.CONNECT_TARGET_HOST)).thenReturn("target_host"); ConnectionDetails cd = cf.getConnectionDetails(PKFactory.Side.TARGET); assertEquals("target_host", cd.host()); } + @Test + public void getConnectionDetailsOriginWithAutoDownloadDisabled() throws Exception { + // Create ConnectionFetcher with mocked AstraDevOpsClient + cf = new ConnectionFetcher(propertyHelper, astraClient); + + // Test + cf.getConnectionDetails(PKFactory.Side.ORIGIN); + + // Verify auto-download was not triggered + verify(astraClient, never()).downloadSecureBundle(any()); + } + + @Test + public void getConnectionDetailsOriginWithAutoDownloadEnabled() throws Exception { + // Setup + when(propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(true); + when(astraClient.downloadSecureBundle(PKFactory.Side.ORIGIN)).thenReturn("/path/to/downloaded/bundle.zip"); + + // Create ConnectionFetcher with mocked AstraDevOpsClient + cf = new ConnectionFetcher(propertyHelper, astraClient); + + // Test + cf.getConnectionDetails(PKFactory.Side.ORIGIN); + + // Verify the SCB path was updated in the property helper + verify(propertyHelper).setProperty(KnownProperties.CONNECT_ORIGIN_SCB, "/path/to/downloaded/bundle.zip"); + } + + @Test + public void getConnectionDetailsTargetWithAutoDownloadEnabled() throws Exception { + // Setup + when(propertyHelper.getBoolean(KnownProperties.TARGET_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(true); + when(astraClient.downloadSecureBundle(PKFactory.Side.TARGET)).thenReturn("/path/to/downloaded/target-bundle.zip"); + + // Create ConnectionFetcher with mocked AstraDevOpsClient + cf = new ConnectionFetcher(propertyHelper, astraClient); + + // Test + cf.getConnectionDetails(PKFactory.Side.TARGET); + + // Verify the SCB path was updated in the property helper + verify(propertyHelper).setProperty(KnownProperties.CONNECT_TARGET_SCB, "/path/to/downloaded/target-bundle.zip"); + } + + @Test + public void getConnectionDetailsWithAutoDownloadFailure() throws Exception { + // Setup + when(propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(true); + when(astraClient.downloadSecureBundle(PKFactory.Side.ORIGIN)).thenThrow(new RuntimeException("Download failed")); + + // Create ConnectionFetcher with mocked AstraDevOpsClient + cf = new ConnectionFetcher(propertyHelper, astraClient); + + // Test - should not throw exception + ConnectionDetails cd = cf.getConnectionDetails(PKFactory.Side.ORIGIN); + + // Verify we still get the connection details + assertEquals("origin_host", cd.host()); + + // And no property update happened + verify(propertyHelper, never()).setProperty(any(), any()); + } } From ec65be758c7ca21a0c73c81eea12cc1b826dd9f1 Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Sat, 3 May 2025 12:16:40 -0400 Subject: [PATCH 2/5] add more unit test coverage --- .../datastax/cdm/data/AstraDevOpsClient.java | 3 +- .../cdm/data/AstraDevOpsClientTest.java | 309 +++++++++++++++++- 2 files changed, 308 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java index 2bb0e144..866ed73e 100644 --- a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java +++ b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java @@ -284,8 +284,7 @@ private String extractDownloadUrl(String jsonResponse, String scbType, PKFactory if (rootNode.has("customDomainBundles") && rootNode.get("customDomainBundles").isArray()) { for (JsonNode customNode : rootNode.get("customDomainBundles")) { - if (customNode.has("sniDomain") - && customDomain.equalsIgnoreCase(customNode.get("sniDomain").asText()) + if (customNode.has("domain") && customDomain.equalsIgnoreCase(customNode.get("domain").asText()) && customNode.has("downloadURL")) { return customNode.get("downloadURL").asText(); diff --git a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java index b74b0202..b8e96a1e 100644 --- a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java +++ b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java @@ -16,17 +16,23 @@ package com.datastax.cdm.data; import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.*; +import java.io.ByteArrayInputStream; import java.io.IOException; +import java.io.InputStream; import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.net.URI; import java.net.http.HttpClient; +import java.net.http.HttpRequest; import java.net.http.HttpResponse; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -46,6 +52,9 @@ class AstraDevOpsClientTest { @Mock private HttpResponse httpResponse; + @Mock + private HttpResponse httpResponseStream; + private AstraDevOpsClient client; @BeforeEach @@ -197,5 +206,301 @@ void testExtractDownloadUrlDefaultType() throws Exception { assertEquals("https://example.com/bundle.zip", url); } - // More tests for extractDownloadUrl with different SCB types can be added here -} \ No newline at end of file + @Test + void testExtractDownloadUrlRegionType() throws Exception { + // Setup + String jsonResponse = "{ \"downloadURLs\": [" + + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," + + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }," + + "{ \"region\": \"eu-central-1\", \"downloadURL\": \"https://eu-central-1.example.com/bundle.zip\" }" + + "]}"; + + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("us-west-2"); + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "region", PKFactory.Side.ORIGIN); + + // Verify + assertEquals("https://us-west-2.example.com/bundle.zip", url); + } + + @Test + void testExtractDownloadUrlRegionTypeWithMissingRegion() throws Exception { + // Setup + String jsonResponse = "{ \"downloadURLs\": [" + + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," + + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }" + "]}"; + + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn(null); + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "region", PKFactory.Side.ORIGIN); + + // Verify + assertNull(url); + } + + @Test + void testExtractDownloadUrlRegionTypeWithNoMatchingRegion() throws Exception { + // Setup + String jsonResponse = "{ \"downloadURLs\": [" + + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," + + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }" + "]}"; + + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("ap-south-1"); + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "region", PKFactory.Side.ORIGIN); + + // Verify + assertNull(url); + } + + @Test + void testExtractDownloadUrlCustomDomainType() throws Exception { + // Setup + String jsonResponse = "{ \"customDomainBundles\": [" + + "{ \"domain\": \"db1.example.com\", \"downloadURL\": \"https://db1.example.com/bundle.zip\" }," + + "{ \"domain\": \"db2.example.com\", \"downloadURL\": \"https://db2.example.com/bundle.zip\" }" + "]}"; + + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN)).thenReturn("db2.example.com"); + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "custom", PKFactory.Side.TARGET); + + // Verify + assertEquals("https://db2.example.com/bundle.zip", url); + } + + @Test + void testExtractDownloadUrlCustomDomainTypeWithMissingDomain() throws Exception { + // Setup + String jsonResponse = "{ \"customDomainBundles\": [" + + "{ \"domain\": \"db1.example.com\", \"downloadURL\": \"https://db1.example.com/bundle.zip\" }," + + "{ \"domain\": \"db2.example.com\", \"downloadURL\": \"https://db2.example.com/bundle.zip\" }" + "]}"; + + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN)).thenReturn(null); + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "custom", PKFactory.Side.TARGET); + + // Verify + assertNull(url); + } + + @Test + void testExtractDownloadUrlUnknownType() throws Exception { + // Setup + String jsonResponse = "{ \"downloadURL\": \"https://example.com/bundle.zip\" }"; + + // Use reflection to access the private method + Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, + String.class, PKFactory.Side.class); + extractDownloadUrlMethod.setAccessible(true); + + // Test + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "unknown", PKFactory.Side.ORIGIN); + + // Verify + assertNull(url); + } + + @Test + void testFetchSecureBundleUrlInfo() throws Exception { + // Mock the HTTP response + when(httpResponse.statusCode()).thenReturn(200); + when(httpResponse.body()).thenReturn("{ \"downloadURL\": \"https://example.com/bundle.zip\" }"); + + // Mock the HTTP client to return our mocked response + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofString()))).thenReturn(httpResponse); + + // Use reflection to access the private method + Method fetchSecureBundleUrlInfoMethod = AstraDevOpsClient.class.getDeclaredMethod( + "fetchSecureBundleUrlInfo", String.class, String.class, boolean.class); + fetchSecureBundleUrlInfoMethod.setAccessible(true); + + // Test + String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "test-token", "test-db-id", false); + + // Verify + assertEquals("{ \"downloadURL\": \"https://example.com/bundle.zip\" }", jsonResponse); + + // Verify the correct URL was used + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClient).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofString())); + + HttpRequest capturedRequest = requestCaptor.getValue(); + assertEquals(URI.create("https://api.astra.datastax.com/v2/databases/test-db-id/secureBundleURL"), + capturedRequest.uri()); + assertTrue(capturedRequest.headers().firstValue("Authorization").isPresent()); + assertEquals("Bearer test-token", capturedRequest.headers().firstValue("Authorization").get()); + } + + @Test + void testFetchSecureBundleUrlInfoWithRegionalFlag() throws Exception { + // Mock the HTTP response + when(httpResponse.statusCode()).thenReturn(200); + when(httpResponse.body()).thenReturn("{ \"downloadURLs\": [] }"); + + // Mock the HTTP client to return our mocked response + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofString()))).thenReturn(httpResponse); + + // Use reflection to access the private method + Method fetchSecureBundleUrlInfoMethod = AstraDevOpsClient.class.getDeclaredMethod( + "fetchSecureBundleUrlInfo", String.class, String.class, boolean.class); + fetchSecureBundleUrlInfoMethod.setAccessible(true); + + // Test + String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "test-token", "test-db-id", true); + + // Verify + assertEquals("{ \"downloadURLs\": [] }", jsonResponse); + + // Verify the correct URL was used with all=true parameter + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClient).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofString())); + + HttpRequest capturedRequest = requestCaptor.getValue(); + assertEquals(URI.create("https://api.astra.datastax.com/v2/databases/test-db-id/secureBundleURL?all=true"), + capturedRequest.uri()); + } + + @Test + void testFetchSecureBundleUrlInfoError() throws Exception { + // Mock the HTTP response for an error + when(httpResponse.statusCode()).thenReturn(401); + when(httpResponse.body()).thenReturn("{ \"error\": \"Unauthorized\" }"); + + // Mock the HTTP client to return our mocked response + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofString()))).thenReturn(httpResponse); + + // Use reflection to access the private method + Method fetchSecureBundleUrlInfoMethod = AstraDevOpsClient.class.getDeclaredMethod( + "fetchSecureBundleUrlInfo", String.class, String.class, boolean.class); + fetchSecureBundleUrlInfoMethod.setAccessible(true); + + // Test + String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "invalid-token", "test-db-id", false); + + // Verify + assertNull(jsonResponse); + } + + @Test + void testDownloadBundleFile() throws Exception { + // Setup + byte[] mockData = new byte[100]; // Mock some binary data + ByteArrayInputStream inputStream = new ByteArrayInputStream(mockData); + + // Mock the HTTP response + when(httpResponseStream.statusCode()).thenReturn(200); + when(httpResponseStream.body()).thenReturn(inputStream); + + // Mock the HTTP client to return our mocked response + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofInputStream()))).thenReturn(httpResponseStream); + + // Use reflection to access the private method + Method downloadBundleFileMethod = AstraDevOpsClient.class.getDeclaredMethod("downloadBundleFile", String.class, + PKFactory.Side.class); + downloadBundleFileMethod.setAccessible(true); + + // Test + String filePath = (String) downloadBundleFileMethod.invoke(client, "https://example.com/bundle.zip", + PKFactory.Side.ORIGIN); + + // Verify + assertNotNull(filePath); + assertTrue(filePath.contains("origin-secure-bundle.zip")); + + // Verify the correct URL was used + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClient).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofInputStream())); + + HttpRequest capturedRequest = requestCaptor.getValue(); + assertEquals(URI.create("https://example.com/bundle.zip"), capturedRequest.uri()); + } + + @Test + void testDownloadBundleFileHttpError() throws Exception { + // Setup + + // Mock the HTTP response with an error status + when(httpResponseStream.statusCode()).thenReturn(404); + + // Mock the HTTP client to return our mocked response + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofInputStream()))).thenReturn(httpResponseStream); + + // Use reflection to access the private method + Method downloadBundleFileMethod = AstraDevOpsClient.class.getDeclaredMethod( + "downloadBundleFile", String.class, PKFactory.Side.class); + downloadBundleFileMethod.setAccessible(true); + + try { + downloadBundleFileMethod.invoke(client, "https://example.com/not-found.zip", PKFactory.Side.ORIGIN); + fail("Expected an exception to be thrown"); + } catch (java.lang.reflect.InvocationTargetException e) { + // Extract the actual exception that was wrapped + assertTrue(e.getCause() instanceof IOException); + assertEquals("Failed to download secure bundle. Status code: 404", e.getCause().getMessage()); + } + } + + @Test + void testDownloadSecureBundleSuccess() throws Exception { + // Setup - mock all the components for a successful download + + // Step 1: Mock the API response for fetching the SCB URL + when(httpResponse.statusCode()).thenReturn(200); + when(httpResponse.body()).thenReturn("{ \"downloadURL\": \"https://example.com/bundle.zip\" }"); + + // Step 2: Mock the binary download + byte[] mockData = new byte[100]; // Mock some binary data + ByteArrayInputStream inputStream = new ByteArrayInputStream(mockData); + + // Mock the HTTP response + when(httpResponseStream.statusCode()).thenReturn(200); + when(httpResponseStream.body()).thenReturn(inputStream); + + // Configure the HTTP client to return responses based on different handler types + // We need to use doReturn().when() syntax here to avoid NullPointerException in the matcher + doReturn(httpResponse).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofString())); + doReturn(httpResponseStream).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofInputStream())); + + // Mock the property helper + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("default"); + + // Test + String filePath = client.downloadSecureBundle(PKFactory.Side.ORIGIN); + + // Verify + assertNotNull(filePath); + assertTrue(filePath.contains("origin-secure-bundle.zip")); + } +} From 6dc08a89e0b57241eb3187938b5c4ca3f81f2e0f Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Sat, 3 May 2025 13:50:07 -0400 Subject: [PATCH 3/5] Reuse CONNECT_PASSWORD as opposed to new Astra token properties --- .../datastax/cdm/data/AstraDevOpsClient.java | 6 +- .../cdm/properties/KnownProperties.java | 1 + src/resources/cdm-detailed.properties | 5 +- .../cdm/data/AstraDevOpsClientTest.java | 174 +++++++++++++++++- 4 files changed, 172 insertions(+), 14 deletions(-) diff --git a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java index 866ed73e..1d29a5e3 100644 --- a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java +++ b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java @@ -113,7 +113,7 @@ public String downloadSecureBundle(PKFactory.Side side) throws IOException { } /** - * Gets the Astra token for the specified side. + * Gets the Astra token for the specified side. Uses the database password as the token for Astra. * * @param side * The side (ORIGIN or TARGET) @@ -121,8 +121,8 @@ public String downloadSecureBundle(PKFactory.Side side) throws IOException { * @return The Astra token */ private String getAstraToken(PKFactory.Side side) { - String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_TOKEN - : KnownProperties.TARGET_ASTRA_TOKEN; + String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.CONNECT_ORIGIN_PASSWORD + : KnownProperties.CONNECT_TARGET_PASSWORD; return propertyHelper.getAsString(property); } diff --git a/src/main/java/com/datastax/cdm/properties/KnownProperties.java b/src/main/java/com/datastax/cdm/properties/KnownProperties.java index 9b5a26af..5e332c1b 100644 --- a/src/main/java/com/datastax/cdm/properties/KnownProperties.java +++ b/src/main/java/com/datastax/cdm/properties/KnownProperties.java @@ -51,6 +51,7 @@ public enum PropertyType { // ========================================================================== // Astra DevOps API Parameters // ========================================================================== + // Note: CONNECT_ORIGIN_PASSWORD and CONNECT_TARGET_PASSWORD properties are used for Astra tokens public static final String ORIGIN_ASTRA_TOKEN = "spark.cdm.connect.origin.astra.token"; public static final String ORIGIN_ASTRA_DATABASE_ID = "spark.cdm.connect.origin.astra.database.id"; public static final String ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB = "spark.cdm.connect.origin.astra.auto.download.scb"; diff --git a/src/resources/cdm-detailed.properties b/src/resources/cdm-detailed.properties index 0ac55c10..06de3547 100644 --- a/src/resources/cdm-detailed.properties +++ b/src/resources/cdm-detailed.properties @@ -65,16 +65,16 @@ spark.cdm.connect.target.password cassandra # # spark.cdm.connect.origin.astra # spark.cdm.connect.target.astra -# .token : The Astra DevOps API token. Required if auto-download is enabled. # .database.id : The Astra Database ID. Required if auto-download is enabled. # .auto.download.scb : Default is false. Set to true to enable auto-download of the SCB. # .scb.type : Default is 'default'. Can be 'default', 'region', or 'custom'. # .scb.region : The region name for regional secure bundles. Required when scb.type is 'region'. # .scb.custom.domain : The custom domain for custom secure bundles. Required when scb.type is 'custom'. +# +# Note: CONNECT_ORIGIN_PASSWORD and CONNECT_TARGET_PASSWORD properties are used for Astra tokens #----------------------------------------------------------------------------------------------------------- # Origin Astra DevOps API configuration -#spark.cdm.connect.origin.astra.token your-astra-devops-api-token #spark.cdm.connect.origin.astra.database.id your-astra-database-id #spark.cdm.connect.origin.astra.auto.download.scb false #spark.cdm.connect.origin.astra.scb.type default @@ -82,7 +82,6 @@ spark.cdm.connect.target.password cassandra #spark.cdm.connect.origin.astra.scb.custom.domain your-custom-domain.example.com # Target Astra DevOps API configuration -#spark.cdm.connect.target.astra.token your-astra-devops-api-token #spark.cdm.connect.target.astra.database.id your-astra-database-id #spark.cdm.connect.target.astra.auto.download.scb false #spark.cdm.connect.target.astra.scb.type default diff --git a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java index b8e96a1e..35a2a50b 100644 --- a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java +++ b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java @@ -71,33 +71,33 @@ void setUp() throws Exception { @Test void testDownloadSecureBundleWithNullToken() throws IOException { // Setup - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn(null); + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn(null); // Test String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); // Verify assertNull(result); - verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN); + verify(propertyHelper).getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD); } @Test void testDownloadSecureBundleWithEmptyToken() throws IOException { // Setup - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn(""); + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn(""); // Test String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); // Verify assertNull(result); - verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN); + verify(propertyHelper).getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD); } @Test void testDownloadSecureBundleWithNullDatabaseId() throws IOException { // Setup - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn(null); // Test @@ -105,14 +105,14 @@ void testDownloadSecureBundleWithNullDatabaseId() throws IOException { // Verify assertNull(result); - verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN); + verify(propertyHelper).getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD); verify(propertyHelper).getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID); } @Test void testGetAstraToken() throws Exception { // Setup - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); // Use reflection to access the private method Method getAstraTokenMethod = AstraDevOpsClient.class.getDeclaredMethod("getAstraToken", PKFactory.Side.class); @@ -492,7 +492,7 @@ void testDownloadSecureBundleSuccess() throws Exception { doReturn(httpResponseStream).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofInputStream())); // Mock the property helper - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_TOKEN)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("default"); @@ -503,4 +503,162 @@ void testDownloadSecureBundleSuccess() throws Exception { assertNotNull(filePath); assertTrue(filePath.contains("origin-secure-bundle.zip")); } + + @Test + void testDownloadBundleFileWithIOException() throws Exception { + // Mock the HTTP client to throw IOException when sending the request + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofInputStream()))) + .thenThrow(new IOException("Network error")); + + // Use reflection to access the private method + Method downloadBundleFileMethod = AstraDevOpsClient.class.getDeclaredMethod( + "downloadBundleFile", String.class, PKFactory.Side.class); + downloadBundleFileMethod.setAccessible(true); + + try { + // Test - should throw an IOException + downloadBundleFileMethod.invoke(client, "https://example.com/bundle.zip", PKFactory.Side.ORIGIN); + fail("Expected an exception to be thrown"); + } catch (java.lang.reflect.InvocationTargetException e) { + // Extract the actual exception that was wrapped + assertTrue(e.getCause() instanceof IOException); + assertEquals("Network error", e.getCause().getMessage()); + } + } + + @Test + void testDownloadSecureBundleWithEmptyJsonResponse() throws Exception { + // Setup + when(httpResponse.statusCode()).thenReturn(200); + when(httpResponse.body()).thenReturn("{}"); + + // Configure the HTTP client to return our mocked response + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofString()))).thenReturn(httpResponse); + + // Mock the property helper + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("default"); + + // Test + String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); + + // Verify + assertNull(result); + } + + @Test + void testDownloadSecureBundleWithRegionalSCB() throws Exception { + // Setup - mock all the components for a successful download with regional SCB + + // Step 1: Mock the API response for fetching the SCB URL with regional data + String jsonResponse = "{ \"downloadURLs\": [" + + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," + + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }," + + "{ \"region\": \"eu-central-1\", \"downloadURL\": \"https://eu-central-1.example.com/bundle.zip\" }" + + "]}"; + + when(httpResponse.statusCode()).thenReturn(200); + when(httpResponse.body()).thenReturn(jsonResponse); + + // Step 2: Mock the binary download + byte[] mockData = new byte[100]; // Mock some binary data + ByteArrayInputStream inputStream = new ByteArrayInputStream(mockData); + + // Mock the HTTP response + when(httpResponseStream.statusCode()).thenReturn(200); + when(httpResponseStream.body()).thenReturn(inputStream); + + // Configure the HTTP client to return responses based on different handler types + doReturn(httpResponse).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofString())); + doReturn(httpResponseStream).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofInputStream())); + + // Mock the property helper + when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("region"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("us-west-2"); + + // Test + String filePath = client.downloadSecureBundle(PKFactory.Side.ORIGIN); + + // Verify + assertNotNull(filePath); + assertTrue(filePath.contains("origin-secure-bundle.zip")); + + // Verify that the URL request included all=true parameter + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClient, atLeastOnce()).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofString())); + + boolean foundAllParam = false; + for (HttpRequest capturedRequest : requestCaptor.getAllValues()) { + if (capturedRequest.uri().toString().contains("all=true")) { + foundAllParam = true; + break; + } + } + assertTrue(foundAllParam, "The request URL should include all=true parameter for regional SCB"); + } + + @Test + void testDownloadSecureBundleWithCustomDomainSCB() throws Exception { + // Setup - mock all the components for a successful download with custom domain SCB + + // Step 1: Mock the API response for fetching the SCB URL with custom domain data + String jsonResponse = "{ \"customDomainBundles\": [" + + "{ \"domain\": \"db1.example.com\", \"downloadURL\": \"https://db1.example.com/bundle.zip\" }," + + "{ \"domain\": \"my-custom-domain.example.com\", \"downloadURL\": \"https://my-custom-domain.example.com/bundle.zip\" }" + + "]}"; + + when(httpResponse.statusCode()).thenReturn(200); + when(httpResponse.body()).thenReturn(jsonResponse); + + // Step 2: Mock the binary download + byte[] mockData = new byte[100]; // Mock some binary data + ByteArrayInputStream inputStream = new ByteArrayInputStream(mockData); + + // Mock the HTTP response + when(httpResponseStream.statusCode()).thenReturn(200); + when(httpResponseStream.body()).thenReturn(inputStream); + + // Configure the HTTP client to return responses based on different handler types + doReturn(httpResponse).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofString())); + doReturn(httpResponseStream).when(httpClient).send(any(), eq(HttpResponse.BodyHandlers.ofInputStream())); + + // Mock the property helper + when(propertyHelper.getAsString(KnownProperties.CONNECT_TARGET_PASSWORD)).thenReturn("test-token"); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_TYPE)).thenReturn("custom"); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN)) + .thenReturn("my-custom-domain.example.com"); + + // Test + String filePath = client.downloadSecureBundle(PKFactory.Side.TARGET); + + // Verify + assertNotNull(filePath); + assertTrue(filePath.contains("target-secure-bundle.zip")); + } + + @Test + void testDownloadBundleFileInterrupted() throws Exception { + // Mock the HTTP client to throw InterruptedException when sending the request + when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofInputStream()))) + .thenThrow(new InterruptedException("Download interrupted")); + + // Use reflection to access the private method + Method downloadBundleFileMethod = AstraDevOpsClient.class.getDeclaredMethod( + "downloadBundleFile", String.class, PKFactory.Side.class); + downloadBundleFileMethod.setAccessible(true); + + try { + // Test - should throw an InterruptedException wrapped in an InvocationTargetException + downloadBundleFileMethod.invoke(client, "https://example.com/bundle.zip", PKFactory.Side.ORIGIN); + fail("Expected an exception to be thrown"); + } catch (java.lang.reflect.InvocationTargetException e) { + // Extract the actual exception that was wrapped + assertTrue(e.getCause() instanceof InterruptedException); + assertEquals("Download interrupted", e.getCause().getMessage()); + } + } } From 024df694692dbe2e03da0991250a5d142ef72edb Mon Sep 17 00:00:00 2001 From: Madhavan Sridharan Date: Mon, 5 May 2025 12:05:40 -0400 Subject: [PATCH 4/5] simplify --- .../datastax/cdm/data/AstraDevOpsClient.java | 93 +++++++----- .../cdm/properties/KnownProperties.java | 10 -- .../datastax/cdm/job/ConnectionFetcher.scala | 13 +- src/resources/cdm-detailed.properties | 14 +- .../cdm/data/AstraDevOpsClientTest.java | 143 ++++++++---------- .../cdm/job/ConnectionFetcherTest.java | 31 ++-- 6 files changed, 153 insertions(+), 151 deletions(-) diff --git a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java index 1d29a5e3..bda4b56a 100644 --- a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java +++ b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java @@ -46,7 +46,7 @@ public class AstraDevOpsClient { private static final Logger logger = LoggerFactory.getLogger(AstraDevOpsClient.class.getName()); private static final String ASTRA_API_BASE_URL = "https://api.astra.datastax.com"; - private static final String SCB_API_PATH = "/v2/databases/%s/secureBundleURL"; + private static final String SCB_API_PATH = "/v2/databases/%s/secureBundleURL?all=true"; private static final Duration HTTP_TIMEOUT = Duration.ofSeconds(30); private final IPropertyHelper propertyHelper; @@ -79,6 +79,7 @@ public AstraDevOpsClient(IPropertyHelper propertyHelper) { public String downloadSecureBundle(PKFactory.Side side) throws IOException { String token = getAstraToken(side); String databaseId = getAstraDatabaseId(side); + String dbRegion = getRegion(side); String scbType = getScbType(side); if (token == null || token.isEmpty() || databaseId == null || databaseId.isEmpty()) { @@ -86,18 +87,18 @@ public String downloadSecureBundle(PKFactory.Side side) throws IOException { return null; } - logger.info("Auto-downloading secure connect bundle for {} database ID: {}, type: {}", side, databaseId, - scbType); + logger.info("Auto-downloading secure connect bundle for {} database ID: {}, type: {}, region: {}", side, + databaseId, scbType, dbRegion != null ? dbRegion : "default"); try { - String jsonResponse = fetchSecureBundleUrlInfo(token, databaseId, "region".equals(scbType)); + String jsonResponse = fetchSecureBundleUrlInfo(token, databaseId); if (jsonResponse == null) { logger.error("Failed to fetch secure bundle URL info for {}", side); return null; } - String downloadUrl = extractDownloadUrl(jsonResponse, scbType, side); + String downloadUrl = extractDownloadUrl(jsonResponse, scbType, side, dbRegion); if (downloadUrl == null) { logger.error("Could not extract download URL for {} bundle type: {}", side, scbType); @@ -135,7 +136,7 @@ private String getAstraToken(PKFactory.Side side) { * * @return The Astra database ID */ - private String getAstraDatabaseId(PKFactory.Side side) { + public String getAstraDatabaseId(PKFactory.Side side) { String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_DATABASE_ID : KnownProperties.TARGET_ASTRA_DATABASE_ID; @@ -165,7 +166,7 @@ private String getScbType(PKFactory.Side side) { * * @return The region name */ - private String getRegion(PKFactory.Side side) { + public String getRegion(PKFactory.Side side) { String property = PKFactory.Side.ORIGIN.equals(side) ? KnownProperties.ORIGIN_ASTRA_SCB_REGION : KnownProperties.TARGET_ASTRA_SCB_REGION; @@ -194,8 +195,6 @@ private String getCustomDomain(PKFactory.Side side) { * The Astra token * @param databaseId * The database ID - * @param fetchAll - * Whether to fetch all regional bundles * * @return The JSON response as a string * @@ -204,15 +203,10 @@ private String getCustomDomain(PKFactory.Side side) { * @throws InterruptedException * If the API call is interrupted */ - private String fetchSecureBundleUrlInfo(String token, String databaseId, boolean fetchAll) - throws IOException, InterruptedException { + private String fetchSecureBundleUrlInfo(String token, String databaseId) throws IOException, InterruptedException { String apiUrl = ASTRA_API_BASE_URL + String.format(SCB_API_PATH, databaseId); - if (fetchAll) { - apiUrl += "?all=true"; - } - HttpRequest request = HttpRequest.newBuilder().uri(URI.create(apiUrl)) .header("Authorization", "Bearer " + token).header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.noBody()).timeout(HTTP_TIMEOUT).build(); @@ -237,41 +231,56 @@ private String fetchSecureBundleUrlInfo(String token, String databaseId, boolean * The SCB type * @param side * The side (ORIGIN or TARGET) + * @param dbRegion + * The database region * * @return The download URL * * @throws IOException * If an error occurs parsing the response */ - private String extractDownloadUrl(String jsonResponse, String scbType, PKFactory.Side side) throws IOException { + private String extractDownloadUrl(String jsonResponse, String scbType, PKFactory.Side side, String dbRegion) + throws IOException { JsonNode rootNode = objectMapper.readTree(jsonResponse); - switch (scbType.toLowerCase()) { - case "default": - // Default bundle URL extraction - if (rootNode.has("downloadURL")) { - return rootNode.get("downloadURL").asText(); + if (!rootNode.isArray()) { + logger.error("Expected array response but got {}", rootNode.getNodeType()); + return null; + } + + // Find the correct datacenter node based on the region + JsonNode matchingDatacenter = null; + + // If region is specified, find the datacenter with that region + if (dbRegion != null && !dbRegion.isEmpty()) { + for (JsonNode datacenterNode : rootNode) { + if (datacenterNode.has("region") && dbRegion.equalsIgnoreCase(datacenterNode.get("region").asText())) { + matchingDatacenter = datacenterNode; + break; + } } - break; - case "region": - // Regional bundle URL extraction - String region = getRegion(side); - if (region == null || region.isEmpty()) { - logger.error("Region is required for SCB type 'region' but was not specified"); + if (matchingDatacenter == null) { + logger.error("Could not find datacenter for region: {}", dbRegion); return null; } + } else { + // No specific region, use the first datacenter + if (rootNode.size() > 0) { + matchingDatacenter = rootNode.get(0); + } else { + logger.error("Response contains no datacenters"); + return null; + } + } - if (rootNode.has("downloadURLs") && rootNode.get("downloadURLs").isArray()) { - for (JsonNode regionNode : rootNode.get("downloadURLs")) { - if (regionNode.has("region") && region.equalsIgnoreCase(regionNode.get("region").asText()) - && regionNode.has("downloadURL")) { - - return regionNode.get("downloadURL").asText(); - } - } - logger.error("Could not find download URL for region: {}", region); + switch (scbType.toLowerCase()) { + case "default": + // Default bundle URL extraction - use the matched datacenter + if (matchingDatacenter.has("downloadURL")) { + return matchingDatacenter.get("downloadURL").asText(); } + logger.error("Could not find default download URL in datacenter"); break; case "custom": @@ -282,15 +291,21 @@ private String extractDownloadUrl(String jsonResponse, String scbType, PKFactory return null; } - if (rootNode.has("customDomainBundles") && rootNode.get("customDomainBundles").isArray()) { - for (JsonNode customNode : rootNode.get("customDomainBundles")) { + if (matchingDatacenter.has("customDomainBundles") + && matchingDatacenter.get("customDomainBundles").isArray()) { + + for (JsonNode customNode : matchingDatacenter.get("customDomainBundles")) { if (customNode.has("domain") && customDomain.equalsIgnoreCase(customNode.get("domain").asText()) && customNode.has("downloadURL")) { return customNode.get("downloadURL").asText(); } } - logger.error("Could not find download URL for custom domain: {}", customDomain); + + logger.error("Could not find downloadURL for custom domain: {} in the selected region {}", customDomain, + dbRegion); + } else { + logger.error("No customDomainBundles found in the selected region {}", dbRegion); } break; diff --git a/src/main/java/com/datastax/cdm/properties/KnownProperties.java b/src/main/java/com/datastax/cdm/properties/KnownProperties.java index 5e332c1b..6632c216 100644 --- a/src/main/java/com/datastax/cdm/properties/KnownProperties.java +++ b/src/main/java/com/datastax/cdm/properties/KnownProperties.java @@ -52,16 +52,12 @@ public enum PropertyType { // Astra DevOps API Parameters // ========================================================================== // Note: CONNECT_ORIGIN_PASSWORD and CONNECT_TARGET_PASSWORD properties are used for Astra tokens - public static final String ORIGIN_ASTRA_TOKEN = "spark.cdm.connect.origin.astra.token"; public static final String ORIGIN_ASTRA_DATABASE_ID = "spark.cdm.connect.origin.astra.database.id"; - public static final String ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB = "spark.cdm.connect.origin.astra.auto.download.scb"; public static final String ORIGIN_ASTRA_SCB_TYPE = "spark.cdm.connect.origin.astra.scb.type"; public static final String ORIGIN_ASTRA_SCB_REGION = "spark.cdm.connect.origin.astra.scb.region"; public static final String ORIGIN_ASTRA_SCB_CUSTOM_DOMAIN = "spark.cdm.connect.origin.astra.scb.custom.domain"; - public static final String TARGET_ASTRA_TOKEN = "spark.cdm.connect.target.astra.token"; public static final String TARGET_ASTRA_DATABASE_ID = "spark.cdm.connect.target.astra.database.id"; - public static final String TARGET_ASTRA_AUTO_DOWNLOAD_SCB = "spark.cdm.connect.target.astra.auto.download.scb"; public static final String TARGET_ASTRA_SCB_TYPE = "spark.cdm.connect.target.astra.scb.type"; public static final String TARGET_ASTRA_SCB_REGION = "spark.cdm.connect.target.astra.scb.region"; public static final String TARGET_ASTRA_SCB_CUSTOM_DOMAIN = "spark.cdm.connect.target.astra.scb.custom.domain"; @@ -88,19 +84,13 @@ public enum PropertyType { defaults.put(CONNECT_TARGET_PASSWORD, "cassandra"); // Astra DevOps API parameters - types.put(ORIGIN_ASTRA_TOKEN, PropertyType.STRING); types.put(ORIGIN_ASTRA_DATABASE_ID, PropertyType.STRING); - types.put(ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB, PropertyType.BOOLEAN); - defaults.put(ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB, "false"); types.put(ORIGIN_ASTRA_SCB_TYPE, PropertyType.STRING); defaults.put(ORIGIN_ASTRA_SCB_TYPE, "default"); types.put(ORIGIN_ASTRA_SCB_REGION, PropertyType.STRING); types.put(ORIGIN_ASTRA_SCB_CUSTOM_DOMAIN, PropertyType.STRING); - types.put(TARGET_ASTRA_TOKEN, PropertyType.STRING); types.put(TARGET_ASTRA_DATABASE_ID, PropertyType.STRING); - types.put(TARGET_ASTRA_AUTO_DOWNLOAD_SCB, PropertyType.BOOLEAN); - defaults.put(TARGET_ASTRA_AUTO_DOWNLOAD_SCB, "false"); types.put(TARGET_ASTRA_SCB_TYPE, PropertyType.STRING); defaults.put(TARGET_ASTRA_SCB_TYPE, "default"); types.put(TARGET_ASTRA_SCB_REGION, PropertyType.STRING); diff --git a/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala b/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala index 4ff3126d..68dcaf54 100644 --- a/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala +++ b/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala @@ -33,17 +33,14 @@ class ConnectionFetcher(propertyHelper: IPropertyHelper, testAstraClient: AstraD } def getConnectionDetails(side: Side): ConnectionDetails = { - // Check if auto-download is enabled for this side - val autoDownloadEnabled = if (Side.ORIGIN.equals(side)) { - propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB) - } else { - propertyHelper.getBoolean(KnownProperties.TARGET_ASTRA_AUTO_DOWNLOAD_SCB) - } + val astraClient = getAstraClient() // If auto-download is enabled, try to download the secure bundle - if (autoDownloadEnabled) { + val astraDbId = astraClient.getAstraDatabaseId(side) + val astraDbRegion = astraClient.getRegion(side) + if (astraDbId != null && !astraDbId.isEmpty && astraDbRegion != null && !astraDbRegion.isEmpty) { + logger.info(s"Auto-downloading secure connect bundle for $side $astraDbId $astraDbRegion") try { - val astraClient = getAstraClient() val downloadedScbPath = astraClient.downloadSecureBundle(side) if (downloadedScbPath != null && !downloadedScbPath.isEmpty) { diff --git a/src/resources/cdm-detailed.properties b/src/resources/cdm-detailed.properties index 06de3547..f40ef396 100644 --- a/src/resources/cdm-detailed.properties +++ b/src/resources/cdm-detailed.properties @@ -60,32 +60,30 @@ spark.cdm.connect.target.password cassandra #=========================================================================================================== # Astra DevOps API Parameters for auto-downloading Secure Connect Bundles +# https://docs.datastax.com/en/astra-db-serverless/databases/secure-connect-bundle.html#download-scbs-with-the-devops-api # These parameters enable automatic download of secure connect bundles from the Astra DevOps API # This eliminates the need to manually download and specify the SCB path # # spark.cdm.connect.origin.astra # spark.cdm.connect.target.astra # .database.id : The Astra Database ID. Required if auto-download is enabled. -# .auto.download.scb : Default is false. Set to true to enable auto-download of the SCB. -# .scb.type : Default is 'default'. Can be 'default', 'region', or 'custom'. # .scb.region : The region name for regional secure bundles. Required when scb.type is 'region'. -# .scb.custom.domain : The custom domain for custom secure bundles. Required when scb.type is 'custom'. +# .scb.type : Default is 'default'. Can be 'default' or 'custom'. +# .scb.custom.domain : The custom domain for secure connect bundles. Required when scb.type is 'custom'. # -# Note: CONNECT_ORIGIN_PASSWORD and CONNECT_TARGET_PASSWORD properties are used for Astra tokens +# Note: 'spark.cdm.connect.origin.password' and 'spark.cdm.connect.target.password' properties are used for Astra tokens #----------------------------------------------------------------------------------------------------------- # Origin Astra DevOps API configuration #spark.cdm.connect.origin.astra.database.id your-astra-database-id -#spark.cdm.connect.origin.astra.auto.download.scb false -#spark.cdm.connect.origin.astra.scb.type default #spark.cdm.connect.origin.astra.scb.region us-east-1 +#spark.cdm.connect.origin.astra.scb.type default #spark.cdm.connect.origin.astra.scb.custom.domain your-custom-domain.example.com # Target Astra DevOps API configuration #spark.cdm.connect.target.astra.database.id your-astra-database-id -#spark.cdm.connect.target.astra.auto.download.scb false -#spark.cdm.connect.target.astra.scb.type default #spark.cdm.connect.target.astra.scb.region us-east-1 +#spark.cdm.connect.target.astra.scb.type default #spark.cdm.connect.target.astra.scb.custom.domain your-custom-domain.example.com #=========================================================================================================== diff --git a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java index 35a2a50b..fcf0c020 100644 --- a/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java +++ b/src/test/java/com/datastax/cdm/data/AstraDevOpsClientTest.java @@ -192,15 +192,16 @@ void testGetCustomDomain() throws Exception { @Test void testExtractDownloadUrlDefaultType() throws Exception { // Setup - String jsonResponse = "{ \"downloadURL\": \"https://example.com/bundle.zip\" }"; + String jsonResponse = "[{ \"downloadURL\": \"https://example.com/bundle.zip\" }]"; // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "default", PKFactory.Side.ORIGIN); + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "default", PKFactory.Side.ORIGIN, + null); // Verify assertEquals("https://example.com/bundle.zip", url); @@ -209,21 +210,22 @@ void testExtractDownloadUrlDefaultType() throws Exception { @Test void testExtractDownloadUrlRegionType() throws Exception { // Setup - String jsonResponse = "{ \"downloadURLs\": [" + String jsonResponse = "[" + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }," + "{ \"region\": \"eu-central-1\", \"downloadURL\": \"https://eu-central-1.example.com/bundle.zip\" }" - + "]}"; + + "]"; - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("us-west-2"); + String region = "us-west-2"; // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "region", PKFactory.Side.ORIGIN); + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "default", PKFactory.Side.ORIGIN, + region); // Verify assertEquals("https://us-west-2.example.com/bundle.zip", url); @@ -232,61 +234,67 @@ void testExtractDownloadUrlRegionType() throws Exception { @Test void testExtractDownloadUrlRegionTypeWithMissingRegion() throws Exception { // Setup - String jsonResponse = "{ \"downloadURLs\": [" + String jsonResponse = "[" + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," - + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }" + "]}"; + + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }" + "]"; - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn(null); + String region = null; // Missing region // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); - // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "region", PKFactory.Side.ORIGIN); + // Test - in new implementation, this should return the first datacenter + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "default", PKFactory.Side.ORIGIN, + region); - // Verify - assertNull(url); + // Verify - we expect it to return the first URL since no region is specified + assertEquals("https://us-east-1.example.com/bundle.zip", url); } @Test void testExtractDownloadUrlRegionTypeWithNoMatchingRegion() throws Exception { // Setup - String jsonResponse = "{ \"downloadURLs\": [" + String jsonResponse = "[" + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," - + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }" + "]}"; + + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }" + "]"; - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("ap-south-1"); + String region = "ap-south-1"; // Non-matching region // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "region", PKFactory.Side.ORIGIN); + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "default", PKFactory.Side.ORIGIN, + region); - // Verify + // Verify - no matching region should return null assertNull(url); } @Test void testExtractDownloadUrlCustomDomainType() throws Exception { // Setup - String jsonResponse = "{ \"customDomainBundles\": [" + String jsonResponse = "[{ " + "\"region\": \"us-east-1\"," + + "\"downloadURL\": \"https://example.com/bundle.zip\"," + "\"customDomainBundles\": [" + "{ \"domain\": \"db1.example.com\", \"downloadURL\": \"https://db1.example.com/bundle.zip\" }," - + "{ \"domain\": \"db2.example.com\", \"downloadURL\": \"https://db2.example.com/bundle.zip\" }" + "]}"; + + "{ \"domain\": \"db2.example.com\", \"downloadURL\": \"https://db2.example.com/bundle.zip\" }" + "]}" + + "]"; when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN)).thenReturn("db2.example.com"); + String region = "us-east-1"; // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "custom", PKFactory.Side.TARGET); + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "custom", PKFactory.Side.TARGET, + region); // Verify assertEquals("https://db2.example.com/bundle.zip", url); @@ -295,36 +303,41 @@ void testExtractDownloadUrlCustomDomainType() throws Exception { @Test void testExtractDownloadUrlCustomDomainTypeWithMissingDomain() throws Exception { // Setup - String jsonResponse = "{ \"customDomainBundles\": [" + String jsonResponse = "[{ " + "\"region\": \"us-east-1\"," + + "\"downloadURL\": \"https://example.com/bundle.zip\"," + "\"customDomainBundles\": [" + "{ \"domain\": \"db1.example.com\", \"downloadURL\": \"https://db1.example.com/bundle.zip\" }," - + "{ \"domain\": \"db2.example.com\", \"downloadURL\": \"https://db2.example.com/bundle.zip\" }" + "]}"; + + "{ \"domain\": \"db2.example.com\", \"downloadURL\": \"https://db2.example.com/bundle.zip\" }" + "]}" + + "]"; when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN)).thenReturn(null); + String region = "us-east-1"; // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); - // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "custom", PKFactory.Side.TARGET); + // Test - missing custom domain + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "custom", PKFactory.Side.TARGET, + region); - // Verify + // Verify - no custom domain should return null assertNull(url); } @Test void testExtractDownloadUrlUnknownType() throws Exception { // Setup - String jsonResponse = "{ \"downloadURL\": \"https://example.com/bundle.zip\" }"; + String jsonResponse = "[{ \"downloadURL\": \"https://example.com/bundle.zip\" }]"; // Use reflection to access the private method Method extractDownloadUrlMethod = AstraDevOpsClient.class.getDeclaredMethod("extractDownloadUrl", String.class, - String.class, PKFactory.Side.class); + String.class, PKFactory.Side.class, String.class); extractDownloadUrlMethod.setAccessible(true); - // Test - String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "unknown", PKFactory.Side.ORIGIN); + // Test with unknown SCB type + String url = (String) extractDownloadUrlMethod.invoke(client, jsonResponse, "unknown", PKFactory.Side.ORIGIN, + null); // Verify assertNull(url); @@ -341,11 +354,11 @@ void testFetchSecureBundleUrlInfo() throws Exception { // Use reflection to access the private method Method fetchSecureBundleUrlInfoMethod = AstraDevOpsClient.class.getDeclaredMethod( - "fetchSecureBundleUrlInfo", String.class, String.class, boolean.class); + "fetchSecureBundleUrlInfo", String.class, String.class); fetchSecureBundleUrlInfoMethod.setAccessible(true); // Test - String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "test-token", "test-db-id", false); + String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "test-token", "test-db-id"); // Verify assertEquals("{ \"downloadURL\": \"https://example.com/bundle.zip\" }", jsonResponse); @@ -355,41 +368,12 @@ void testFetchSecureBundleUrlInfo() throws Exception { verify(httpClient).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofString())); HttpRequest capturedRequest = requestCaptor.getValue(); - assertEquals(URI.create("https://api.astra.datastax.com/v2/databases/test-db-id/secureBundleURL"), + assertEquals(URI.create("https://api.astra.datastax.com/v2/databases/test-db-id/secureBundleURL?all=true"), capturedRequest.uri()); assertTrue(capturedRequest.headers().firstValue("Authorization").isPresent()); assertEquals("Bearer test-token", capturedRequest.headers().firstValue("Authorization").get()); } - @Test - void testFetchSecureBundleUrlInfoWithRegionalFlag() throws Exception { - // Mock the HTTP response - when(httpResponse.statusCode()).thenReturn(200); - when(httpResponse.body()).thenReturn("{ \"downloadURLs\": [] }"); - - // Mock the HTTP client to return our mocked response - when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofString()))).thenReturn(httpResponse); - - // Use reflection to access the private method - Method fetchSecureBundleUrlInfoMethod = AstraDevOpsClient.class.getDeclaredMethod( - "fetchSecureBundleUrlInfo", String.class, String.class, boolean.class); - fetchSecureBundleUrlInfoMethod.setAccessible(true); - - // Test - String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "test-token", "test-db-id", true); - - // Verify - assertEquals("{ \"downloadURLs\": [] }", jsonResponse); - - // Verify the correct URL was used with all=true parameter - ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); - verify(httpClient).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofString())); - - HttpRequest capturedRequest = requestCaptor.getValue(); - assertEquals(URI.create("https://api.astra.datastax.com/v2/databases/test-db-id/secureBundleURL?all=true"), - capturedRequest.uri()); - } - @Test void testFetchSecureBundleUrlInfoError() throws Exception { // Mock the HTTP response for an error @@ -401,11 +385,11 @@ void testFetchSecureBundleUrlInfoError() throws Exception { // Use reflection to access the private method Method fetchSecureBundleUrlInfoMethod = AstraDevOpsClient.class.getDeclaredMethod( - "fetchSecureBundleUrlInfo", String.class, String.class, boolean.class); + "fetchSecureBundleUrlInfo", String.class, String.class); fetchSecureBundleUrlInfoMethod.setAccessible(true); // Test - String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "invalid-token", "test-db-id", false); + String jsonResponse = (String) fetchSecureBundleUrlInfoMethod.invoke(client, "invalid-token", "test-db-id"); // Verify assertNull(jsonResponse); @@ -476,7 +460,7 @@ void testDownloadSecureBundleSuccess() throws Exception { // Step 1: Mock the API response for fetching the SCB URL when(httpResponse.statusCode()).thenReturn(200); - when(httpResponse.body()).thenReturn("{ \"downloadURL\": \"https://example.com/bundle.zip\" }"); + when(httpResponse.body()).thenReturn("[{ \"downloadURL\": \"https://example.com/bundle.zip\" }]"); // Step 2: Mock the binary download byte[] mockData = new byte[100]; // Mock some binary data @@ -495,6 +479,7 @@ void testDownloadSecureBundleSuccess() throws Exception { when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("default"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn(null); // Test String filePath = client.downloadSecureBundle(PKFactory.Side.ORIGIN); @@ -530,7 +515,7 @@ void testDownloadBundleFileWithIOException() throws Exception { void testDownloadSecureBundleWithEmptyJsonResponse() throws Exception { // Setup when(httpResponse.statusCode()).thenReturn(200); - when(httpResponse.body()).thenReturn("{}"); + when(httpResponse.body()).thenReturn("[]"); // Empty array response // Configure the HTTP client to return our mocked response when(httpClient.send(any(), eq(HttpResponse.BodyHandlers.ofString()))).thenReturn(httpResponse); @@ -539,6 +524,7 @@ void testDownloadSecureBundleWithEmptyJsonResponse() throws Exception { when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("default"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn(null); // Test String result = client.downloadSecureBundle(PKFactory.Side.ORIGIN); @@ -552,11 +538,11 @@ void testDownloadSecureBundleWithRegionalSCB() throws Exception { // Setup - mock all the components for a successful download with regional SCB // Step 1: Mock the API response for fetching the SCB URL with regional data - String jsonResponse = "{ \"downloadURLs\": [" + String jsonResponse = "[" + "{ \"region\": \"us-east-1\", \"downloadURL\": \"https://us-east-1.example.com/bundle.zip\" }," + "{ \"region\": \"us-west-2\", \"downloadURL\": \"https://us-west-2.example.com/bundle.zip\" }," + "{ \"region\": \"eu-central-1\", \"downloadURL\": \"https://eu-central-1.example.com/bundle.zip\" }" - + "]}"; + + "]"; when(httpResponse.statusCode()).thenReturn(200); when(httpResponse.body()).thenReturn(jsonResponse); @@ -576,7 +562,7 @@ void testDownloadSecureBundleWithRegionalSCB() throws Exception { // Mock the property helper when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_PASSWORD)).thenReturn("test-token"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("test-db-id"); - when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("region"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_TYPE)).thenReturn("default"); when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("us-west-2"); // Test @@ -590,6 +576,7 @@ void testDownloadSecureBundleWithRegionalSCB() throws Exception { ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpRequest.class); verify(httpClient, atLeastOnce()).send(requestCaptor.capture(), eq(HttpResponse.BodyHandlers.ofString())); + // Verify the URL includes all=true parameter boolean foundAllParam = false; for (HttpRequest capturedRequest : requestCaptor.getAllValues()) { if (capturedRequest.uri().toString().contains("all=true")) { @@ -605,10 +592,11 @@ void testDownloadSecureBundleWithCustomDomainSCB() throws Exception { // Setup - mock all the components for a successful download with custom domain SCB // Step 1: Mock the API response for fetching the SCB URL with custom domain data - String jsonResponse = "{ \"customDomainBundles\": [" + String jsonResponse = "[{ " + "\"region\": \"us-east-1\"," + + "\"downloadURL\": \"https://example.com/bundle.zip\"," + "\"customDomainBundles\": [" + "{ \"domain\": \"db1.example.com\", \"downloadURL\": \"https://db1.example.com/bundle.zip\" }," + "{ \"domain\": \"my-custom-domain.example.com\", \"downloadURL\": \"https://my-custom-domain.example.com/bundle.zip\" }" - + "]}"; + + "]}" + "]"; when(httpResponse.statusCode()).thenReturn(200); when(httpResponse.body()).thenReturn(jsonResponse); @@ -631,6 +619,7 @@ void testDownloadSecureBundleWithCustomDomainSCB() throws Exception { when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_TYPE)).thenReturn("custom"); when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_CUSTOM_DOMAIN)) .thenReturn("my-custom-domain.example.com"); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_REGION)).thenReturn("us-east-1"); // Test String filePath = client.downloadSecureBundle(PKFactory.Side.TARGET); diff --git a/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java b/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java index fbe0bced..a82b4069 100644 --- a/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java +++ b/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java @@ -63,9 +63,11 @@ public void setup() { when(propertyHelper.getAsString(KnownProperties.CONNECT_ORIGIN_HOST)).thenReturn("origin_host"); when(propertyHelper.getAsString(KnownProperties.CONNECT_TARGET_HOST)).thenReturn("target_host"); - // Default - auto-download disabled for both ORIGIN and TARGET - when(propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(false); - when(propertyHelper.getBoolean(KnownProperties.TARGET_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(false); + // Default - auto-download disabled for both ORIGIN and TARGET by not setting database ID and region + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn(null); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn(null); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_DATABASE_ID)).thenReturn(null); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_REGION)).thenReturn(null); } @Test @@ -94,8 +96,11 @@ public void getConnectionDetailsOriginWithAutoDownloadDisabled() throws Exceptio @Test public void getConnectionDetailsOriginWithAutoDownloadEnabled() throws Exception { - // Setup - when(propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(true); + // Setup - auto-download enabled by setting both database ID and region + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("origin-db-id"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("us-east-1"); + when(astraClient.getAstraDatabaseId(PKFactory.Side.ORIGIN)).thenReturn("origin-db-id"); + when(astraClient.getRegion(PKFactory.Side.ORIGIN)).thenReturn("us-east-1"); when(astraClient.downloadSecureBundle(PKFactory.Side.ORIGIN)).thenReturn("/path/to/downloaded/bundle.zip"); // Create ConnectionFetcher with mocked AstraDevOpsClient @@ -110,8 +115,11 @@ public void getConnectionDetailsOriginWithAutoDownloadEnabled() throws Exception @Test public void getConnectionDetailsTargetWithAutoDownloadEnabled() throws Exception { - // Setup - when(propertyHelper.getBoolean(KnownProperties.TARGET_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(true); + // Setup - auto-download enabled by setting both database ID and region + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_DATABASE_ID)).thenReturn("target-db-id"); + when(propertyHelper.getAsString(KnownProperties.TARGET_ASTRA_SCB_REGION)).thenReturn("eu-west-1"); + when(astraClient.getAstraDatabaseId(PKFactory.Side.TARGET)).thenReturn("target-db-id"); + when(astraClient.getRegion(PKFactory.Side.TARGET)).thenReturn("eu-west-1"); when(astraClient.downloadSecureBundle(PKFactory.Side.TARGET)).thenReturn("/path/to/downloaded/target-bundle.zip"); // Create ConnectionFetcher with mocked AstraDevOpsClient @@ -126,8 +134,13 @@ public void getConnectionDetailsTargetWithAutoDownloadEnabled() throws Exception @Test public void getConnectionDetailsWithAutoDownloadFailure() throws Exception { - // Setup - when(propertyHelper.getBoolean(KnownProperties.ORIGIN_ASTRA_AUTO_DOWNLOAD_SCB)).thenReturn(true); + // Setup - auto-download enabled by setting both database ID and region + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_DATABASE_ID)).thenReturn("origin-db-id"); + when(propertyHelper.getAsString(KnownProperties.ORIGIN_ASTRA_SCB_REGION)).thenReturn("us-east-1"); + when(astraClient.getAstraDatabaseId(PKFactory.Side.ORIGIN)).thenReturn("origin-db-id"); + when(astraClient.getRegion(PKFactory.Side.ORIGIN)).thenReturn("us-east-1"); + + // But the download fails when(astraClient.downloadSecureBundle(PKFactory.Side.ORIGIN)).thenThrow(new RuntimeException("Download failed")); // Create ConnectionFetcher with mocked AstraDevOpsClient From ca792860ecc1e4779de893611440d621696b93bf Mon Sep 17 00:00:00 2001 From: Pravin Bhat Date: Tue, 6 May 2025 09:42:10 -0400 Subject: [PATCH 5/5] File path fix --- src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java | 6 ------ src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala | 2 +- .../java/com/datastax/cdm/job/ConnectionFetcherTest.java | 6 ++---- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java index bda4b56a..aba6ec23 100644 --- a/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java +++ b/src/main/java/com/datastax/cdm/data/AstraDevOpsClient.java @@ -15,7 +15,6 @@ */ package com.datastax.cdm.data; -import java.io.File; import java.io.FileOutputStream; import java.io.IOException; import java.io.InputStream; @@ -25,10 +24,7 @@ import java.net.http.HttpResponse; import java.nio.file.Files; import java.nio.file.Path; -import java.nio.file.Paths; import java.time.Duration; -import java.util.Map; -import java.util.Optional; import java.util.UUID; import org.slf4j.Logger; @@ -204,9 +200,7 @@ private String getCustomDomain(PKFactory.Side side) { * If the API call is interrupted */ private String fetchSecureBundleUrlInfo(String token, String databaseId) throws IOException, InterruptedException { - String apiUrl = ASTRA_API_BASE_URL + String.format(SCB_API_PATH, databaseId); - HttpRequest request = HttpRequest.newBuilder().uri(URI.create(apiUrl)) .header("Authorization", "Bearer " + token).header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.noBody()).timeout(HTTP_TIMEOUT).build(); diff --git a/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala b/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala index 68dcaf54..0d773372 100644 --- a/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala +++ b/src/main/scala/com/datastax/cdm/job/ConnectionFetcher.scala @@ -41,7 +41,7 @@ class ConnectionFetcher(propertyHelper: IPropertyHelper, testAstraClient: AstraD if (astraDbId != null && !astraDbId.isEmpty && astraDbRegion != null && !astraDbRegion.isEmpty) { logger.info(s"Auto-downloading secure connect bundle for $side $astraDbId $astraDbRegion") try { - val downloadedScbPath = astraClient.downloadSecureBundle(side) + val downloadedScbPath = "file://" + astraClient.downloadSecureBundle(side) if (downloadedScbPath != null && !downloadedScbPath.isEmpty) { logger.info(s"Successfully auto-downloaded secure bundle for $side: $downloadedScbPath") diff --git a/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java b/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java index a82b4069..ee4b7d07 100644 --- a/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java +++ b/src/test/java/com/datastax/cdm/job/ConnectionFetcherTest.java @@ -21,8 +21,6 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import java.lang.reflect.Field; - import org.apache.spark.SparkConf; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -110,7 +108,7 @@ public void getConnectionDetailsOriginWithAutoDownloadEnabled() throws Exception cf.getConnectionDetails(PKFactory.Side.ORIGIN); // Verify the SCB path was updated in the property helper - verify(propertyHelper).setProperty(KnownProperties.CONNECT_ORIGIN_SCB, "/path/to/downloaded/bundle.zip"); + verify(propertyHelper).setProperty(KnownProperties.CONNECT_ORIGIN_SCB, "file:///path/to/downloaded/bundle.zip"); } @Test @@ -129,7 +127,7 @@ public void getConnectionDetailsTargetWithAutoDownloadEnabled() throws Exception cf.getConnectionDetails(PKFactory.Side.TARGET); // Verify the SCB path was updated in the property helper - verify(propertyHelper).setProperty(KnownProperties.CONNECT_TARGET_SCB, "/path/to/downloaded/target-bundle.zip"); + verify(propertyHelper).setProperty(KnownProperties.CONNECT_TARGET_SCB, "file:///path/to/downloaded/target-bundle.zip"); } @Test