From fa19dc3dae1eb1e416943243a9314e9b1913cec3 Mon Sep 17 00:00:00 2001 From: andreadimaio Date: Mon, 28 Sep 2026 13:41:22 +0200 Subject: [PATCH] Replace COS service singleton with per-project-ID cache --- .../text-extraction-service.md | 2 +- .../ai/textprocessing/CosDataLocation.java | 13 +- .../storage/cos/CosStorageService.java | 9 +- .../storage/cos/DefaultRestClient.java | 10 +- .../TextClassificationService.java | 154 ++++++++-------- .../textextraction/TextExtractionService.java | 173 +++++++++--------- .../TextClassificationTest.java | 71 +++++++ .../textextraction/TextExtractionTest.java | 81 ++++++++ 8 files changed, 342 insertions(+), 171 deletions(-) diff --git a/docs/content/services/document-processing/text-extraction-service.md b/docs/content/services/document-processing/text-extraction-service.md index f12aad0c..d84aa415 100644 --- a/docs/content/services/document-processing/text-extraction-service.md +++ b/docs/content/services/document-processing/text-extraction-service.md @@ -113,7 +113,7 @@ TextExtractionService service = TextExtractionService.builder() .build(); ``` -The resolved endpoint URL and bucket are cached after the first lookup. When only `spaceId` is configured (no `projectId`), lazy COS resolution is not supported; set `cosUrl` explicitly in that case. +The resolved endpoint URL and bucket are cached per Project ID after the first lookup for that project. When only `spaceId` is configured (no `projectId`), lazy COS resolution is not supported. Set `cosUrl` explicitly in that case. ### Using a Separate COS Authenticator diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/CosDataLocation.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/CosDataLocation.java index 0ac90883..68162ffc 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/CosDataLocation.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/CosDataLocation.java @@ -15,4 +15,15 @@ * @param bucket Optional override for the bucket name defined in the connection asset. Used only for {@code connection_asset} type. * @param path The file path within the container. Used only for {@code container} type. */ -public record CosDataLocation(String fileName, String bucket, String path) {} +public record CosDataLocation(String fileName, String bucket, String path) { + + public CosDataLocation { + boolean hasConnectionAssetFields = fileName != null || bucket != null; + boolean hasPath = path != null; + if (hasConnectionAssetFields == hasPath) + throw new IllegalArgumentException( + hasPath + ? "fileName/bucket and path are mutually exclusive" + : "Either fileName/bucket or path must be set"); + } +} diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/CosStorageService.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/CosStorageService.java index d51ad157..6d42df49 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/CosStorageService.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/CosStorageService.java @@ -18,7 +18,6 @@ import com.ibm.watsonx.ai.core.auth.ibmcloud.IBMCloudAuthenticator; import com.ibm.watsonx.ai.core.exception.WatsonxException; import com.ibm.watsonx.ai.core.exception.model.WatsonxError.Code; -import com.ibm.watsonx.ai.core.exception.model.WatsonxError.Error; import com.ibm.watsonx.ai.core.http.logging.HttpRequestLogger; import com.ibm.watsonx.ai.core.http.logging.HttpResponseLogger; import com.ibm.watsonx.ai.textprocessing.DeleteFileRequest; @@ -108,6 +107,7 @@ public boolean deleteFile(String requestId, String fileName) throws FileNotFound try { return client.deleteFileAsync(DeleteFileRequest.of(requestId, null, fileName)).get(); } catch (InterruptedException e) { + Thread.currentThread().interrupt(); throw new RuntimeException(e); } catch (ExecutionException e) { if (e.getCause() instanceof WatsonxException ex) { @@ -352,12 +352,11 @@ public CosStorageService build() { private Exception mapIfCosFileNotFound(WatsonxException e) { if (e.statusCode() == 404 && e.details().isPresent()) { var details = e.details().get(); - var fileNotFound = details.errors().stream() + return details.errors().stream() .filter(error -> error.is(Code.COS_FILE_NOT_FOUND)) .findFirst() - .map(Error::message) - .orElse(e.getMessage()); - return new FileNotFoundException(fileNotFound); + .map(error -> new FileNotFoundException(error.message())) + .orElse(e); } return e; } diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/DefaultRestClient.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/DefaultRestClient.java index 9c58ae39..4860988a 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/DefaultRestClient.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/storage/cos/DefaultRestClient.java @@ -71,7 +71,10 @@ public boolean upload(String requestTrackingId, InputStream is, String fileName) var response = syncCosHttpClient.send(builder.build(), BodyHandlers.ofString()); return response.statusCode() == 200; - } catch (IOException | InterruptedException | URISyntaxException e) { + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } catch (IOException | URISyntaxException e) { throw new RuntimeException(e); } } @@ -102,7 +105,10 @@ public String readFile(String requestTrackingId, String fileName) throws FileNot builder.header(REQUEST_ID_HEADER, requestTrackingId); return syncCosHttpClient.send(builder.build(), BodyHandlers.ofString()).body(); - } catch (IOException | InterruptedException | URISyntaxException e) { + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } catch (IOException | URISyntaxException e) { throw new RuntimeException(e); } } diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationService.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationService.java index 60a591de..31de7678 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationService.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationService.java @@ -18,7 +18,7 @@ import java.time.Duration; import java.util.Map; import java.util.UUID; -import java.util.concurrent.locks.ReentrantLock; +import java.util.concurrent.ConcurrentHashMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import com.ibm.watsonx.ai.CloudRegion; @@ -64,8 +64,7 @@ public class TextClassificationService extends ScopedService { private final String cosUrl; private final DocumentReference documentReference; private final TextClassificationRestClient client; - private volatile StorageOperations cosService; - private final ReentrantLock cosServiceLock = new ReentrantLock(); + private final ConcurrentHashMap cosServiceCache = new ConcurrentHashMap<>(); private final ProjectService lazyProjectService; private final Authenticator authenticator; private final Authenticator cosAuthenticator; @@ -379,8 +378,7 @@ public boolean uploadFile(InputStream inputStream, String fileName) { requireNonNull(inputStream, "inputStream value cannot be null"); requireNonNull(fileName, "fileName value cannot be null"); var requestId = UUID.randomUUID().toString(); - upload(requestId, inputStream, fileName, null, false); - return true; + return upload(requestId, inputStream, fileName, null, false); } /** @@ -442,81 +440,81 @@ public boolean deleteRequest(String id, TextClassificationDeleteParameters param return client.deleteClassification(request); } - // Returns the CosStorageService lazily on first use. private StorageOperations getOrResolveCosService() { - return getOrResolveCosService(null); + return getOrResolveCosService(this.projectId); } private StorageOperations getOrResolveCosService(CosReference cosRef) { + return StorageFactory.cos(cosUrl, cosRef.bucket(), authenticator, cosAuthenticator, httpClient, + timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); + } - if (cosRef != null) { - // Per-call CosReference override. - return StorageFactory.cos(cosUrl, cosRef.bucket(), authenticator, cosAuthenticator, httpClient, - timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); - } - - if (cosService != null) - return cosService; + private StorageOperations getOrResolveCosService(String effectiveProjectId) { + if (effectiveProjectId == null) + throw new IllegalStateException( + "ContainerReference storage requires a projectId - spaceId alone is not supported for container uploads."); - cosServiceLock.lock(); + var cached = cosServiceCache.get(effectiveProjectId); + if (cached != null) + return cached; + boolean wasInterrupted = Thread.interrupted(); try { + return cosServiceCache.computeIfAbsent(effectiveProjectId, this::resolveCosServiceForProject); + } finally { + if (wasInterrupted) + Thread.currentThread().interrupt(); + } + } - if (cosService != null) - return cosService; - - if (projectId == null) - throw new IllegalStateException( - "ContainerReference storage requires a projectId - spaceId alone is not supported for container uploads."); - - ProjectService ps = lazyProjectService; - - if (ps == null) { + private String resolveEffectiveProjectId(TextClassificationParameters parameters) { + if (parameters != null && (parameters.projectId() != null || parameters.spaceId() != null)) + return parameters.projectId(); + return this.projectId; + } - var region = CloudRegion.fromMlEndpoint(baseUrl).orElseThrow(() -> new IllegalStateException( - "ContainerReference storage requires a ProjectService or a known CloudRegion. " - + "Either pass projectService(ProjectService) on the builder, " - + "or use baseUrl(CloudRegion) so the project storage can be resolved automatically.")); + private StorageOperations resolveCosServiceForProject(String projectId) { + ProjectService ps = lazyProjectService; - var psBuilder = ProjectService.builder() - .baseUrl(region.wxEndpoint().replace("/wx", "")) - .authenticator(authenticator) - .httpClient(httpClient) - .timeout(timeout) - .logRequests(logRequests) - .logResponses(logResponses) - .verifySsl(verifySsl); + if (ps == null) { - if (requestLogger != null) - psBuilder.logRequests(requestLogger, requestLogLevel); + var region = CloudRegion.fromMlEndpoint(baseUrl).orElseThrow(() -> new IllegalStateException( + "ContainerReference storage requires a ProjectService or a known CloudRegion. " + + "Either pass projectService(ProjectService) on the builder, " + + "or use baseUrl(CloudRegion) so the project storage can be resolved automatically.")); - if (responseLogger != null) - psBuilder.logResponses(responseLogger, responseLogLevel); + var psBuilder = ProjectService.builder() + .baseUrl(region.wxEndpoint().replace("/wx", "")) + .authenticator(authenticator) + .httpClient(httpClient) + .timeout(timeout) + .logRequests(logRequests) + .logResponses(logResponses) + .verifySsl(verifySsl); - ps = psBuilder.build(); - } + if (requestLogger != null) + psBuilder.logRequests(requestLogger, requestLogLevel); - var props = ps.findProject(projectId) - .orElseThrow(() -> new IllegalStateException("Project not found: " + projectId)) - .storage() - .properties(); + if (responseLogger != null) + psBuilder.logResponses(responseLogger, responseLogLevel); - var resolvedUrl = props.endpointUrl().endsWith("/") - ? props.endpointUrl().substring(0, props.endpointUrl().length() - 1) - : props.endpointUrl(); + ps = psBuilder.build(); + } - var resolvedBucket = props.bucketName(); - cosService = StorageFactory.cos(resolvedUrl, resolvedBucket, authenticator, cosAuthenticator, httpClient, - timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); + var props = ps.findProject(projectId) + .orElseThrow(() -> new IllegalStateException("Project not found: " + projectId)) + .storage() + .properties(); - return cosService; + var resolvedUrl = props.endpointUrl().endsWith("/") + ? props.endpointUrl().substring(0, props.endpointUrl().length() - 1) + : props.endpointUrl(); - } finally { - cosServiceLock.unlock(); - } + var resolvedBucket = props.bucketName(); + return StorageFactory.cos(resolvedUrl, resolvedBucket, authenticator, cosAuthenticator, httpClient, + timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); } - // Retrieves the classification result for the given request id. private TextClassificationResponse fetchClassificationRequest(String requestId, String id, TextClassificationFetchParameters parameters) { requireNonNull(requestId, "The requestId can not be null"); requireNonNull(id, "The id can not be null"); @@ -550,13 +548,12 @@ private ClassificationResult classifyAndFetch(String requestId, String absoluteP DocumentReference effectiveDoc = parameters.documentReference() != null ? parameters.documentReference() : this.documentReference; - cleanUpUploadedFile(requestId, uploadedPath, effectiveDoc); + cleanUpUploadedFile(requestId, uploadedPath, effectiveDoc, resolveEffectiveProjectId(parameters)); } } } - // Uploads an input stream to COS or the container. - private void upload(String requestId, InputStream is, String fileName, TextClassificationParameters parameters, + private boolean upload(String requestId, InputStream is, String fileName, TextClassificationParameters parameters, boolean waitForClassification) { requireNonNull(requestId, "requestId value cannot be null"); requireNonNull(is, "is value cannot be null"); @@ -573,10 +570,8 @@ private void upload(String requestId, InputStream is, String fileName, TextClass this.documentReference ); - if (effectiveDoc instanceof ContainerReference) { - getOrResolveCosService().upload(requestId, is, fileName); - return; - } + if (effectiveDoc instanceof ContainerReference) + return getOrResolveCosService(resolveEffectiveProjectId(parameters)).upload(requestId, is, fileName); if (!(effectiveDoc instanceof CosReference cosRef)) throw new UnsupportedOperationException( @@ -586,7 +581,7 @@ private void upload(String requestId, InputStream is, String fileName, TextClass throw new IllegalStateException( "cosUrl must be set on the service builder to perform COS upload operations."); - getOrResolveCosService(cosRef).upload(requestId, is, fileName); + return getOrResolveCosService(cosRef).upload(requestId, is, fileName); } // Starts the text classification process. @@ -705,8 +700,7 @@ private void cleanUpAfterAbortedClassification(String processId, String projectI } } - // Deletes the uploaded file from COS or the container. All errors are logged and swallowed. - private void cleanUpUploadedFile(String requestId, String path, DocumentReference effectiveDoc) { + private void cleanUpUploadedFile(String requestId, String path, DocumentReference effectiveDoc, String effectiveProjectId) { try { if (effectiveDoc instanceof CosReference cosRef) client.deleteFileAsync(DeleteFileRequest.of(requestId, cosRef.bucket(), path)) @@ -714,21 +708,13 @@ private void cleanUpUploadedFile(String requestId, String path, DocumentReferenc logger.warn("Async COS delete failed for {}: {}", path, ex.getMessage()); return false; }); - else if (effectiveDoc instanceof ContainerReference) { - boolean wasInterrupted = Thread.interrupted(); - try { - getOrResolveCosService().deleteFileAsync(requestId, path); - } finally { - if (wasInterrupted) - Thread.currentThread().interrupt(); - } - } + else if (effectiveDoc instanceof ContainerReference) + getOrResolveCosService(effectiveProjectId).deleteFileAsync(requestId, path); } catch (Exception e) { logger.warn("Failed to delete uploaded file {}: {}", path, e.getMessage(), e); } } - // Extracts the ClassificationResult from a completed-or-failed response. private ClassificationResult getClassificationResult(TextClassificationResponse textClassificationResponse) throws TextClassificationException { @@ -834,6 +820,16 @@ public Builder documentReference(ContainerReference documentReference) { return this; } + /** + * Specifies the reference to the input document, whether a {@link CosReference} or a {@link ContainerReference}. + * + * @param documentReference the {@link DocumentReference} for the input file. + */ + public Builder documentReference(DocumentReference documentReference) { + this.documentReference = documentReference; + return this; + } + /** * Specifies the Cloud Object Storage connection and bucket where the input files are stored. * diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionService.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionService.java index 01963dbd..b006f05c 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionService.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionService.java @@ -21,7 +21,7 @@ import java.util.List; import java.util.Map; import java.util.UUID; -import java.util.concurrent.locks.ReentrantLock; +import java.util.concurrent.ConcurrentHashMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import com.ibm.watsonx.ai.CloudRegion; @@ -70,8 +70,7 @@ public class TextExtractionService extends ScopedService { private final DocumentReference documentReference; private final DocumentReference resultReference; private final TextExtractionRestClient client; - private volatile StorageOperations cosService; - private final ReentrantLock cosServiceLock = new ReentrantLock(); + private final ConcurrentHashMap cosServiceCache = new ConcurrentHashMap<>(); private final ProjectService lazyProjectService; private final Authenticator authenticator; private final Authenticator cosAuthenticator; @@ -408,8 +407,7 @@ public boolean uploadFile(InputStream inputStream, String fileName) { requireNonNull(inputStream, "inputStream value cannot be null"); requireNonNull(fileName, "fileName value cannot be null"); var requestId = UUID.randomUUID().toString(); - upload(requestId, inputStream, fileName, null, false); - return true; + return upload(requestId, inputStream, fileName, null, false); } /** @@ -491,75 +489,77 @@ public boolean deleteRequest(String id, TextExtractionDeleteParameters parameter return client.deleteExtraction(request); } - // Returns the CosStorageService lazily on first use. private StorageOperations getOrResolveCosService() { - return getOrResolveCosService(null); + return getOrResolveCosService(this.projectId); } private StorageOperations getOrResolveCosService(CosReference cosRef) { + return StorageFactory.cos(cosUrl, cosRef.bucket(), authenticator, cosAuthenticator, httpClient, + timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); + } - if (cosRef != null) { - // Per-call CosReference override. - return StorageFactory.cos(cosUrl, cosRef.bucket(), authenticator, cosAuthenticator, httpClient, - timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); - } - - if (cosService != null) - return cosService; + private StorageOperations getOrResolveCosService(String effectiveProjectId) { + if (effectiveProjectId == null) + throw new IllegalStateException( + "ContainerReference storage requires a projectId - spaceId alone is not supported for container uploads."); - cosServiceLock.lock(); + var cached = cosServiceCache.get(effectiveProjectId); + if (cached != null) + return cached; + boolean wasInterrupted = Thread.interrupted(); try { + return cosServiceCache.computeIfAbsent(effectiveProjectId, this::resolveCosServiceForProject); + } finally { + if (wasInterrupted) + Thread.currentThread().interrupt(); + } + } - if (cosService != null) - return cosService; - - if (projectId == null) - throw new IllegalStateException( - "ContainerReference storage requires a projectId - spaceId alone is not supported for container uploads."); - - ProjectService ps = lazyProjectService; - - if (ps == null) { - var region = CloudRegion.fromMlEndpoint(baseUrl).orElseThrow(() -> new IllegalStateException( - "ContainerReference storage requires a ProjectService or a known CloudRegion. " - + "Either pass projectService(ProjectService) on the builder, " - + "or use baseUrl(CloudRegion) so the project storage can be resolved automatically.")); - var psBuilder = ProjectService.builder() - .baseUrl(region.wxEndpoint().replace("/wx", "")) - .authenticator(authenticator) - .httpClient(httpClient) - .timeout(timeout) - .logRequests(logRequests) - .logResponses(logResponses) - .verifySsl(verifySsl); - - if (requestLogger != null) - psBuilder.logRequests(requestLogger, requestLogLevel); - - if (responseLogger != null) - psBuilder.logResponses(responseLogger, responseLogLevel); - - ps = psBuilder.build(); - } + private String resolveEffectiveProjectId(TextExtractionParameters parameters) { + if (parameters != null && (parameters.projectId() != null || parameters.spaceId() != null)) + return parameters.projectId(); + return this.projectId; + } - var props = ps.findProject(projectId) - .orElseThrow(() -> new IllegalStateException("Project not found: " + projectId)) - .storage().properties(); + private StorageOperations resolveCosServiceForProject(String projectId) { + ProjectService ps = lazyProjectService; + + if (ps == null) { + var region = CloudRegion.fromMlEndpoint(baseUrl).orElseThrow(() -> new IllegalStateException( + "ContainerReference storage requires a ProjectService or a known CloudRegion. " + + "Either pass projectService(ProjectService) on the builder, " + + "or use baseUrl(CloudRegion) so the project storage can be resolved automatically.")); + var psBuilder = ProjectService.builder() + .baseUrl(region.wxEndpoint().replace("/wx", "")) + .authenticator(authenticator) + .httpClient(httpClient) + .timeout(timeout) + .logRequests(logRequests) + .logResponses(logResponses) + .verifySsl(verifySsl); + + if (requestLogger != null) + psBuilder.logRequests(requestLogger, requestLogLevel); + + if (responseLogger != null) + psBuilder.logResponses(responseLogger, responseLogLevel); + + ps = psBuilder.build(); + } - var resolvedUrl = props.endpointUrl().endsWith("/") - ? props.endpointUrl().substring(0, props.endpointUrl().length() - 1) - : props.endpointUrl(); + var props = ps.findProject(projectId) + .orElseThrow(() -> new IllegalStateException("Project not found: " + projectId)) + .storage().properties(); - var resolvedBucket = props.bucketName(); + var resolvedUrl = props.endpointUrl().endsWith("/") + ? props.endpointUrl().substring(0, props.endpointUrl().length() - 1) + : props.endpointUrl(); - cosService = StorageFactory.cos(resolvedUrl, resolvedBucket, authenticator, cosAuthenticator, httpClient, - timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); + var resolvedBucket = props.bucketName(); - return cosService; - } finally { - cosServiceLock.unlock(); - } + return StorageFactory.cos(resolvedUrl, resolvedBucket, authenticator, cosAuthenticator, httpClient, + timeout, logRequests, logResponses, requestLogger, requestLogLevel, responseLogger, responseLogLevel, verifySsl); } // Retrieves the results of a text extraction request by its unique identifier. @@ -598,12 +598,11 @@ private String extractAndFetch(String requestId, String absolutePath, TextExtrac DocumentReference effectiveDoc = parameters.documentReference() != null ? parameters.documentReference() : this.documentReference; - cleanUpUploadedFile(requestId, uploadedPath, effectiveDoc); + cleanUpUploadedFile(requestId, uploadedPath, effectiveDoc, resolveEffectiveProjectId(parameters)); } } } - // Validates that the requested outputs are compatible with a single-file fetch operation. private void validateFetchOutputs(TextExtractionParameters parameters) throws TextExtractionException { if (parameters == null) @@ -621,8 +620,7 @@ private void validateFetchOutputs(TextExtractionParameters parameters) throws Te "The fetch operation cannot be executed for the type \"page_images\""); } - // Uploads an input stream to COS or the container. - private void upload(String requestId, InputStream is, String fileName, TextExtractionParameters parameters, + private boolean upload(String requestId, InputStream is, String fileName, TextExtractionParameters parameters, boolean waitForExtraction) { requireNonNull(requestId, "requestId value cannot be null"); requireNonNull(is, "is value cannot be null"); @@ -639,10 +637,8 @@ private void upload(String requestId, InputStream is, String fileName, TextExtra this.documentReference ); - if (effectiveDoc instanceof ContainerReference) { - getOrResolveCosService().upload(requestId, is, fileName); - return; - } + if (effectiveDoc instanceof ContainerReference) + return getOrResolveCosService(resolveEffectiveProjectId(parameters)).upload(requestId, is, fileName); if (!(effectiveDoc instanceof CosReference cosRef)) throw new UnsupportedOperationException( @@ -651,7 +647,7 @@ private void upload(String requestId, InputStream is, String fileName, TextExtra throw new IllegalStateException( "cosUrl must be set on the service builder to perform COS upload operations."); - getOrResolveCosService(cosRef).upload(requestId, is, fileName); + return getOrResolveCosService(cosRef).upload(requestId, is, fileName); } // Starts the text extraction process. @@ -800,8 +796,7 @@ private void cleanUpAfterAbortedExtraction(String processId, String projectId, S } } - // Deletes the uploaded file from COS or the container. All errors are logged and swallowed. - private void cleanUpUploadedFile(String requestId, String path, DocumentReference effectiveDoc) { + private void cleanUpUploadedFile(String requestId, String path, DocumentReference effectiveDoc, String effectiveProjectId) { try { if (effectiveDoc instanceof CosReference cosRef) client.deleteFileAsync(DeleteFileRequest.of(requestId, cosRef.bucket(), path)) @@ -809,21 +804,13 @@ private void cleanUpUploadedFile(String requestId, String path, DocumentReferenc logger.warn("Async COS delete failed for {}: {}", path, ex.getMessage()); return false; }); - else if (effectiveDoc instanceof ContainerReference) { - boolean wasInterrupted = Thread.interrupted(); - try { - getOrResolveCosService().deleteFileAsync(requestId, path); - } finally { - if (wasInterrupted) - Thread.currentThread().interrupt(); - } - } + else if (effectiveDoc instanceof ContainerReference) + getOrResolveCosService(effectiveProjectId).deleteFileAsync(requestId, path); } catch (Exception e) { logger.warn("Failed to delete uploaded file {}: {}", path, e.getMessage(), e); } } - // Retrieves the extracted text from COS or the container. private String getExtractedText(String requestId, TextExtractionResponse textExtractionResponse, TextExtractionParameters parameters) throws TextExtractionException, FileNotFoundException { @@ -845,13 +832,13 @@ private String getExtractedText(String requestId, TextExtractionResponse textExt String outputPath = resultsLocation.path() != null ? resultsLocation.path() : resultsLocation.fileName(); String extractedFile; if (isContainerResult) - extractedFile = getOrResolveCosService().readFile(requestId, outputPath); + extractedFile = getOrResolveCosService(resolveEffectiveProjectId(parameters)).readFile(requestId, outputPath); else extractedFile = client.readFile(ReadFileRequest.of(requestId, resultsBucketName, outputPath)); if (removeOutputFile) { try { if (isContainerResult) { - getOrResolveCosService().deleteFileAsync(requestId, outputPath); + getOrResolveCosService(resolveEffectiveProjectId(parameters)).deleteFileAsync(requestId, outputPath); } else { client.deleteFileAsync(DeleteFileRequest.of(requestId, resultsBucketName, outputPath)) .exceptionally(ex -> { @@ -967,6 +954,16 @@ public Builder documentReference(ContainerReference documentReference) { return this; } + /** + * Specifies the reference to the input document, whether a {@link CosReference} or a {@link ContainerReference}. + * + * @param documentReference the {@link DocumentReference} for the input file. + */ + public Builder documentReference(DocumentReference documentReference) { + this.documentReference = documentReference; + return this; + } + /** * Specifies the Cloud Object Storage connection and bucket where the input files are stored. * @@ -1001,6 +998,16 @@ public Builder resultReference(ContainerReference resultReference) { return this; } + /** + * Specifies the reference to the output location, whether a {@link CosReference} or a {@link ContainerReference}. + * + * @param resultReference the {@link DocumentReference} for the output location. + */ + public Builder resultReference(DocumentReference resultReference) { + this.resultReference = resultReference; + return this; + } + /** * Specifies the Cloud Object Storage connection and bucket where the extracted results should be stored. * diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationTest.java index 14e543e6..1695c678 100644 --- a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationTest.java +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textclassification/TextClassificationTest.java @@ -1015,6 +1015,17 @@ void should_upload_file() throws Exception { assertTrue(classificationService.uploadFile(file)); } + @Test + void should_return_false_when_cos_upload_does_not_respond_200() throws Exception { + + var file = new File(ClassLoader.getSystemResource("test.pdf").toURI()); + cosServer.stubFor(put("/%s/%s".formatted("my-bucket", "test.pdf")) + .withHeader("Authorization", equalTo("Bearer token")) + .willReturn(aResponse().withStatus(201))); + + assertFalse(classificationService.uploadFile(file)); + } + @Test @MockitoSettings(strictness = Strictness.LENIENT) void should_upload_file_with_different_api_key() throws Exception { @@ -1645,6 +1656,66 @@ void should_upload_to_project_bucket_when_per_call_container_reference_overrides cosServer.verify(0, putRequestedFor(urlPathMatching("/cos-bucket/.*"))); } + @Test + void should_resolve_separate_cos_service_per_project_when_parameters_override_project_id() throws Exception { + + when(mockAuthenticator.token()).thenReturn("token"); + + var mockProjectService = mock(ProjectService.class); + + var mockProjectA = mock(Project.class); + var mockStorageA = mock(ProjectStorage.class); + var mockPropsA = mock(ProjectStorageProperties.class); + when(mockProjectService.findProject("project-id")).thenReturn(Optional.of(mockProjectA)); + when(mockProjectA.storage()).thenReturn(mockStorageA); + when(mockStorageA.properties()).thenReturn(mockPropsA); + when(mockPropsA.endpointUrl()).thenReturn("http://localhost:%s".formatted(cosServer.getPort())); + when(mockPropsA.bucketName()).thenReturn("bucket-a"); + + var mockProjectB = mock(Project.class); + var mockStorageB = mock(ProjectStorage.class); + var mockPropsB = mock(ProjectStorageProperties.class); + when(mockProjectService.findProject("project-b")).thenReturn(Optional.of(mockProjectB)); + when(mockProjectB.storage()).thenReturn(mockStorageB); + when(mockStorageB.properties()).thenReturn(mockPropsB); + when(mockPropsB.endpointUrl()).thenReturn("http://localhost:%s".formatted(cosServer.getPort())); + when(mockPropsB.bucketName()).thenReturn("bucket-b"); + + var JOB = Files.readString(Path.of(ClassLoader.getSystemResource("classification_job.json").toURI())); + + var service = TextClassificationService.builder() + .baseUrl("http://localhost:%s".formatted(watsonxServer.getPort())) + .authenticator(mockAuthenticator) + .projectId("project-id") + .documentReference(ContainerReference.container()) + .projectService(mockProjectService) + .build(); + + cosServer.stubFor(put(urlPathMatching("/bucket-a/.*")).willReturn(aResponse().withStatus(200))); + cosServer.stubFor(put(urlPathMatching("/bucket-b/.*")).willReturn(aResponse().withStatus(200))); + watsonxServer.stubFor(post(urlPathEqualTo("/ml/v1/text/classifications")) + .willReturn(aResponse().withStatus(200).withHeader("Content-Type", "application/json").withBody(JOB.formatted("submitted")))); + + var file = new File(ClassLoader.getSystemResource("test.pdf").toURI()); + + // Default project ("project-id") resolves and caches bucket-a. + service.uploadAndStartClassification(file); + cosServer.verify(1, putRequestedFor(urlPathMatching("/bucket-a/.*"))); + + // A per-call override to "project-b" must resolve bucket-b, not reuse the cached bucket-a. + var overrideParams = TextClassificationParameters.builder().projectId("project-b").build(); + service.uploadAndStartClassification(file, overrideParams); + cosServer.verify(1, putRequestedFor(urlPathMatching("/bucket-b/.*"))); + + // Back on the default project, the cache must still resolve bucket-a, not the bucket cached for "project-b". + service.uploadAndStartClassification(file); + cosServer.verify(2, putRequestedFor(urlPathMatching("/bucket-a/.*"))); + cosServer.verify(1, putRequestedFor(urlPathMatching("/bucket-b/.*"))); + + verify(mockProjectService, times(1)).findProject("project-id"); + verify(mockProjectService, times(1)).findProject("project-b"); + } + @Test void should_return_result_when_delete_of_uploaded_file_responds_403() throws Exception { diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionTest.java index 9aea67ca..b72a23dc 100644 --- a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionTest.java +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textprocessing/textextraction/TextExtractionTest.java @@ -2231,6 +2231,18 @@ void should_upload_file_with_different_api_key() throws Exception { assertTrue(textExtractionService.uploadFile(file)); } + @Test + void should_return_false_when_cos_upload_does_not_respond_200() throws Exception { + + when(mockAuthenticator.token()).thenReturn("my-super-token"); + var file = new File(ClassLoader.getSystemResource("test.pdf").toURI()); + cosServer.stubFor(put("/%s/%s".formatted("my-bucket", "test.pdf")) + .withHeader("Authorization", equalTo("Bearer my-super-token")) + .willReturn(aResponse().withStatus(201))); + + assertFalse(textExtractionService.uploadFile(file)); + } + @Test @MockitoSettings(strictness = Strictness.LENIENT) void should_throw_exception_when_uploading_non_existent_file_to_cos() throws Exception { @@ -2521,6 +2533,75 @@ void should_upload_to_project_bucket_when_per_call_container_reference_overrides cosServer.verify(0, putRequestedFor(urlPathMatching("/cos-bucket/.*"))); } + @Test + void should_resolve_separate_cos_service_per_project_when_parameters_override_project_id() throws Exception { + + when(mockAuthenticator.token()).thenReturn("my-super-token"); + + var mockProjectService = mock(ProjectService.class); + + var mockProjectA = mock(Project.class); + var mockStorageA = mock(ProjectStorage.class); + var mockPropsA = mock(ProjectStorageProperties.class); + when(mockProjectService.findProject("projectid")).thenReturn(Optional.of(mockProjectA)); + when(mockProjectA.storage()).thenReturn(mockStorageA); + when(mockStorageA.properties()).thenReturn(mockPropsA); + when(mockPropsA.endpointUrl()).thenReturn("http://localhost:%s".formatted(cosServer.getPort())); + when(mockPropsA.bucketName()).thenReturn("bucket-a"); + + var mockProjectB = mock(Project.class); + var mockStorageB = mock(ProjectStorage.class); + var mockPropsB = mock(ProjectStorageProperties.class); + when(mockProjectService.findProject("project-b")).thenReturn(Optional.of(mockProjectB)); + when(mockProjectB.storage()).thenReturn(mockStorageB); + when(mockStorageB.properties()).thenReturn(mockPropsB); + when(mockPropsB.endpointUrl()).thenReturn("http://localhost:%s".formatted(cosServer.getPort())); + when(mockPropsB.bucketName()).thenReturn("bucket-b"); + + var RESPONSE = """ + { + "metadata": { "id": "%s", "created_at": "2023-05-02T16:27:51Z", "project_id": "projectid" }, + "entity": { + "document_reference": { "type": "container", "location": { "path": "test.pdf" } }, + "results_reference": { "type": "container", "location": { "path": "test.md" } }, + "results": { "status": "submitted", "number_pages_processed": 0 } + } + }""".formatted(PROCESS_EXTRACTION_ID); + + var service = TextExtractionService.builder() + .baseUrl("http://localhost:%s".formatted(watsonxServer.getPort())) + .authenticator(mockAuthenticator) + .projectId("projectid") + .documentReference(ContainerReference.container()) + .resultReference(ContainerReference.container()) + .projectService(mockProjectService) + .build(); + + cosServer.stubFor(put(urlPathMatching("/bucket-a/.*")).willReturn(aResponse().withStatus(200))); + cosServer.stubFor(put(urlPathMatching("/bucket-b/.*")).willReturn(aResponse().withStatus(200))); + watsonxServer.stubFor(post(urlPathEqualTo("/ml/v1/text/extractions")) + .willReturn(aResponse().withStatus(200).withHeader("Content-Type", "application/json").withBody(RESPONSE))); + + var file = new File(ClassLoader.getSystemResource(FILE_NAME).toURI()); + + // Default project ("projectid") resolves and caches bucket-a. + service.uploadAndStartExtraction(file); + cosServer.verify(1, putRequestedFor(urlPathMatching("/bucket-a/.*"))); + + // A per-call override to "project-b" must resolve bucket-b, not reuse the cached bucket-a. + var overrideParams = TextExtractionParameters.builder().projectId("project-b").build(); + service.uploadAndStartExtraction(file, overrideParams); + cosServer.verify(1, putRequestedFor(urlPathMatching("/bucket-b/.*"))); + + // Back on the default project, the cache must still resolve bucket-a, not the bucket cached for "project-b". + service.uploadAndStartExtraction(file); + cosServer.verify(2, putRequestedFor(urlPathMatching("/bucket-a/.*"))); + cosServer.verify(1, putRequestedFor(urlPathMatching("/bucket-b/.*"))); + + verify(mockProjectService, times(1)).findProject("projectid"); + verify(mockProjectService, times(1)).findProject("project-b"); + } + @Test void should_still_delete_uploaded_file_when_output_file_not_found() throws Exception {