This is an automated email from the ASF dual-hosted git repository.
pvillard31 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new af81f18baab NIFI-15911 Add Web Client Service for Snowflake Ingest
Service (#11369)
af81f18baab is described below
commit af81f18baabdeb987eafd344648a15d9ba60c936
Author: David Handermann <[email protected]>
AuthorDate: Thu Jun 25 09:40:45 2026 -0500
NIFI-15911 Add Web Client Service for Snowflake Ingest Service (#11369)
---
.../nifi-snowflake-services/pom.xml | 14 +++
.../snowflake/service/SnowpipeIngestClient.java | 115 ++++++++++-----------
...ndardSnowflakeIngestManagerProviderService.java | 36 ++++---
.../service/SnowpipeIngestClientTest.java | 8 +-
...ndardSnowflakeIngestManagerProviderService.java | 25 ++++-
5 files changed, 122 insertions(+), 76 deletions(-)
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
index 1a2dd5e1e97..d6884454068 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/pom.xml
@@ -69,6 +69,14 @@
<groupId>org.apache.nifi</groupId>
<artifactId>nifi-oauth2-provider-api</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-web-client-provider-api</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-web-client-api</artifactId>
+ </dependency>
<dependency>
<groupId>org.apache.nifi</groupId>
@@ -89,5 +97,11 @@
<artifactId>mockwebserver3-junit5</artifactId>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.nifi</groupId>
+ <artifactId>nifi-web-client</artifactId>
+ <version>2.11.0-SNAPSHOT</version>
+ <scope>test</scope>
+ </dependency>
</dependencies>
</project>
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowpipeIngestClient.java
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowpipeIngestClient.java
index c9406556869..5e805a65bd9 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowpipeIngestClient.java
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/SnowpipeIngestClient.java
@@ -22,23 +22,21 @@ import
com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.nifi.processors.snowflake.snowpipe.InsertFiles;
import org.apache.nifi.processors.snowflake.snowpipe.InsertReport;
+import org.apache.nifi.web.client.api.HttpResponseEntity;
+import org.apache.nifi.web.client.api.WebClientService;
+import org.apache.nifi.web.client.api.WebClientServiceException;
import java.io.IOException;
+import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URI;
-import java.net.http.HttpClient;
-import java.net.http.HttpRequest;
-import java.net.http.HttpResponse;
-import java.time.Duration;
import java.util.Objects;
import java.util.UUID;
/**
* Client for Snowflake Snowpipe REST API using Java HttpClient
*/
-class SnowpipeIngestClient implements AutoCloseable {
-
- private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(10);
+class SnowpipeIngestClient {
private static final String INSERT_FILES_PATH =
"/v1/data/pipes/%s/insertFiles";
@@ -56,7 +54,7 @@ class SnowpipeIngestClient implements AutoCloseable {
private static final ObjectMapper objectMapper = new
ObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
- private final HttpClient httpClient;
+ private final WebClientService webClientService;
private final URI insertFilesUri;
@@ -74,11 +72,13 @@ class SnowpipeIngestClient implements AutoCloseable {
SnowpipeIngestClient(
final URI baseUri,
final String pipeName,
- final RSAKeyAuthorizationProvider authorizationProvider
+ final RSAKeyAuthorizationProvider authorizationProvider,
+ final WebClientService webClientService
) {
Objects.requireNonNull(baseUri, "Base URI required");
Objects.requireNonNull(pipeName, "Pipe Name required");
Objects.requireNonNull(authorizationProvider, "Authorization Provider
required");
+ Objects.requireNonNull(webClientService, "Web Service Client
required");
final String insertFilesPath = INSERT_FILES_PATH.formatted(pipeName);
this.insertFilesUri = baseUri.resolve(insertFilesPath);
@@ -87,14 +87,7 @@ class SnowpipeIngestClient implements AutoCloseable {
this.insertReportUri = baseUri.resolve(insertReportPath);
this.authorizationProvider = authorizationProvider;
- this.httpClient = HttpClient.newBuilder()
- .connectTimeout(CONNECT_TIMEOUT)
- .build();
- }
-
- @Override
- public void close() {
- httpClient.close();
+ this.webClientService = webClientService;
}
/**
@@ -110,18 +103,25 @@ class SnowpipeIngestClient implements AutoCloseable {
final URI requestUri = appendRequestId(insertFilesUri);
final String authorization =
authorizationProvider.getRequestAuthorization().authorization();
- final HttpRequest request = HttpRequest.newBuilder()
- .uri(requestUri)
- .header(AUTHORIZATION_HEADER, authorization)
- .header(CONTENT_TYPE_HEADER, APPLICATION_JSON)
- .POST(HttpRequest.BodyPublishers.ofString(requestBody))
- .build();
-
- final HttpResponse<String> response = sendRequest(request);
- final int statusCode = response.statusCode();
- if (statusCode != HttpURLConnection.HTTP_OK) {
- final String message = "Insert Files POST [%s] HTTP %d Response
[%s]".formatted(requestUri, statusCode, response.body());
- throw new SnowpipeResponseException(message);
+ try (
+ HttpResponseEntity responseEntity = webClientService.post()
+ .uri(requestUri)
+ .header(AUTHORIZATION_HEADER, authorization)
+ .header(CONTENT_TYPE_HEADER, APPLICATION_JSON)
+ .body(requestBody)
+ .retrieve();
+ InputStream responseBodyStream = responseEntity.body()
+ ) {
+ final int statusCode = responseEntity.statusCode();
+ if (statusCode != HttpURLConnection.HTTP_OK) {
+ final byte[] responseBodyBytes =
responseBodyStream.readAllBytes();
+ final String responseBody = new String(responseBodyBytes);
+ final String message = "Insert Files POST [%s] HTTP %d
Response [%s]".formatted(requestUri, statusCode, responseBody);
+ throw new SnowpipeResponseException(message);
+ }
+ } catch (final IOException | WebClientServiceException e) {
+ final String message = "Insert Files POST [%s] request
failed".formatted(requestUri);
+ throw new SnowpipeResponseException(message, e);
}
}
@@ -135,26 +135,34 @@ class SnowpipeIngestClient implements AutoCloseable {
final URI requestUri = appendRequestId(insertReportUri);
final String authorization =
authorizationProvider.getRequestAuthorization().authorization();
- final HttpRequest request = HttpRequest.newBuilder()
- .uri(requestUri)
- .header(AUTHORIZATION_HEADER, authorization)
- .header(ACCEPT_HEADER, APPLICATION_JSON)
- .GET()
- .build();
-
- final HttpResponse<String> response = sendRequest(request);
- final int statusCode = response.statusCode();
- final String responseBody = response.body();
- if (statusCode == HttpURLConnection.HTTP_OK) {
- try {
- return objectMapper.readValue(responseBody,
InsertReport.class);
- } catch (final JsonProcessingException e) {
- final String message = "Insert Report [%s] response parsing
failed [%s]".formatted(requestUri, responseBody);
- throw new SnowpipeResponseException(message, e);
+ try (
+ HttpResponseEntity responseEntity = webClientService.get()
+ .uri(requestUri)
+ .header(AUTHORIZATION_HEADER, authorization)
+ .header(ACCEPT_HEADER, APPLICATION_JSON)
+ .retrieve();
+ InputStream responseBodyStream = responseEntity.body()
+ ) {
+ final int statusCode = responseEntity.statusCode();
+
+ final byte[] responseBodyBytes = responseBodyStream.readAllBytes();
+ final String responseBody = new String(responseBodyBytes);
+
+ if (statusCode == HttpURLConnection.HTTP_OK) {
+ try {
+ return objectMapper.readValue(responseBody,
InsertReport.class);
+ } catch (final JsonProcessingException e) {
+ final String message = "Insert Report [%s] response
parsing failed [%s]".formatted(requestUri, responseBody);
+ throw new SnowpipeResponseException(message, e);
+ }
+ } else {
+ final String message = "Insert Report failed GET [%s] HTTP %d
Response [%s]".formatted(requestUri, statusCode, responseBody);
+ throw new SnowpipeResponseException(message);
}
- } else {
- final String message = "Insert Report failed GET [%s] HTTP %d
Response [%s]".formatted(requestUri, statusCode, responseBody);
- throw new SnowpipeResponseException(message);
+
+ } catch (final IOException | WebClientServiceException e) {
+ final String message = "Insert Report [%s] request
failed".formatted(requestUri);
+ throw new SnowpipeResponseException(message, e);
}
}
@@ -166,17 +174,6 @@ class SnowpipeIngestClient implements AutoCloseable {
}
}
- private HttpResponse<String> sendRequest(final HttpRequest request) {
- try {
- return httpClient.send(request,
HttpResponse.BodyHandlers.ofString());
- } catch (final IOException e) {
- throw new SnowpipeResponseException("Request failed for URI
[%s]".formatted(request.uri()), e);
- } catch (final InterruptedException e) {
- Thread.currentThread().interrupt();
- throw new SnowpipeResponseException("Request interrupted for URI
[%s]".formatted(request.uri()), e);
- }
- }
-
private URI appendRequestId(final URI baseUri) {
final String requestId = UUID.randomUUID().toString();
final String uriString = REQUEST_ID_FORMAT.formatted(baseUri,
requestId);
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/StandardSnowflakeIngestManagerProviderService.java
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/StandardSnowflakeIngestManagerProviderService.java
index d7d0b049f86..a703cff2928 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/StandardSnowflakeIngestManagerProviderService.java
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/main/java/org/apache/nifi/snowflake/service/StandardSnowflakeIngestManagerProviderService.java
@@ -19,7 +19,6 @@ package org.apache.nifi.snowflake.service;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
-import org.apache.nifi.annotation.lifecycle.OnDisabled;
import org.apache.nifi.annotation.lifecycle.OnEnabled;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.controller.AbstractControllerService;
@@ -36,11 +35,15 @@ import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.snowflake.service.util.AccountIdentifierFormat;
import
org.apache.nifi.snowflake.service.util.AccountIdentifierFormatParameters;
import org.apache.nifi.snowflake.service.util.ConnectionUrlFormat;
+import org.apache.nifi.web.client.api.WebClientService;
+import org.apache.nifi.web.client.provider.api.WebClientServiceProvider;
import java.net.URI;
import java.security.PrivateKey;
import java.security.interfaces.RSAPrivateCrtKey;
import java.util.List;
+import java.util.Map;
+import java.util.Optional;
@Tags({"snowflake", "snowpipe", "database", "connection"})
@CapabilityDescription("Provides a Snowflake Ingest Manager for Snowflake pipe
processors")
@@ -106,6 +109,13 @@ public class StandardSnowflakeIngestManagerProviderService
extends AbstractContr
.required(true)
.build();
+ public static final PropertyDescriptor WEB_CLIENT_SERVICE_PROVIDER = new
PropertyDescriptor.Builder()
+ .name("Web Client Service Provider")
+ .description("Provider implementation for Web Client Service
supporting HTTP request and response handling")
+ .identifiesControllerService(WebClientServiceProvider.class)
+ .required(true)
+ .build();
+
public static final PropertyDescriptor DATABASE = new
PropertyDescriptor.Builder()
.fromPropertyDescriptor(SnowflakeProperties.DATABASE)
.required(true)
@@ -134,6 +144,7 @@ public class StandardSnowflakeIngestManagerProviderService
extends AbstractContr
ACCOUNT_NAME,
USER_NAME,
PRIVATE_KEY_SERVICE,
+ WEB_CLIENT_SERVICE_PROVIDER,
DATABASE,
SCHEMA,
PIPE
@@ -147,13 +158,13 @@ public class
StandardSnowflakeIngestManagerProviderService extends AbstractContr
private static final char HYPHEN = '-';
+ private volatile SnowpipeIngestClient ingestClient;
+
@Override
protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
return PROPERTY_DESCRIPTORS;
}
- private volatile SnowpipeIngestClient ingestClient;
-
@OnEnabled
public void onEnabled(final ConfigurationContext context) throws
InitializationException {
final String user =
context.getProperty(USER_NAME).evaluateAttributeExpressions().getValue();
@@ -173,20 +184,15 @@ public class
StandardSnowflakeIngestManagerProviderService extends AbstractContr
final URI baseUri =
URI.create(HTTPS_URI_FORMAT.formatted(hostNormalized));
final RSAKeyAuthorizationProvider authorizationProvider = new
RSAKeyAuthorizationProvider(account, user, rsaPrivateKey);
- ingestClient = new SnowpipeIngestClient(baseUri,
qualifiedPipeName, authorizationProvider);
+
+ final WebClientServiceProvider webClientServiceProvider =
context.getProperty(WEB_CLIENT_SERVICE_PROVIDER).asControllerService(WebClientServiceProvider.class);
+ final WebClientService webClientService =
webClientServiceProvider.getWebClientService();
+ ingestClient = new SnowpipeIngestClient(baseUri,
qualifiedPipeName, authorizationProvider, webClientService);
} else {
throw new InitializationException("RSA Private Key not provided");
}
}
- @OnDisabled
- public void onDisabled() {
- if (ingestClient != null) {
- ingestClient.close();
- ingestClient = null;
- }
- }
-
@Override
public void insertFiles(final InsertFiles insertFiles) {
ingestClient.insertFiles(insertFiles);
@@ -211,6 +217,12 @@ public class StandardSnowflakeIngestManagerProviderService
extends AbstractContr
config.renameProperty(SnowflakeProperties.OLD_ACCOUNT_NAME_PROPERTY_NAME,
SnowflakeProperties.ACCOUNT_NAME.getName());
config.renameProperty(SnowflakeProperties.OLD_DATABASE_PROPERTY_NAME,
SnowflakeProperties.DATABASE.getName());
config.renameProperty(SnowflakeProperties.OLD_SCHEMA_PROPERTY_NAME,
SnowflakeProperties.SCHEMA.getName());
+
+ final Optional<String> webClientServiceProviderConfigured =
config.getPropertyValue(WEB_CLIENT_SERVICE_PROVIDER);
+ if (webClientServiceProviderConfigured.isEmpty()) {
+ final String webClientServiceProviderId =
config.createControllerService("org.apache.nifi.web.client.provider.service.StandardWebClientServiceProvider",
Map.of());
+ config.setProperty(WEB_CLIENT_SERVICE_PROVIDER,
webClientServiceProviderId);
+ }
}
private AccountIdentifierFormatParameters
getAccountIdentifierFormatParameters(ConfigurationContext context) {
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowpipeIngestClientTest.java
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowpipeIngestClientTest.java
index 57cb5dd082b..bc01314d763 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowpipeIngestClientTest.java
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/SnowpipeIngestClientTest.java
@@ -26,6 +26,7 @@ import
org.apache.nifi.processors.snowflake.snowpipe.InsertFile;
import org.apache.nifi.processors.snowflake.snowpipe.InsertFileStatus;
import org.apache.nifi.processors.snowflake.snowpipe.InsertFiles;
import org.apache.nifi.processors.snowflake.snowpipe.InsertReport;
+import org.apache.nifi.web.client.StandardWebClientService;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -128,6 +129,8 @@ class SnowpipeIngestClientTest {
@StartStop
public final MockWebServer mockWebServer = new MockWebServer();
+ private StandardWebClientService webClientService;
+
private SnowpipeIngestClient client;
@BeforeEach
@@ -135,12 +138,13 @@ class SnowpipeIngestClientTest {
final URI baseUri =
URI.create(HTTP_URI_FORMAT.formatted(mockWebServer.getHostName(),
mockWebServer.getPort()));
final RSAPrivateCrtKey privateKey = generatePrivateKey();
final RSAKeyAuthorizationProvider authProvider = new
RSAKeyAuthorizationProvider(ACCOUNT, USER, privateKey);
- client = new SnowpipeIngestClient(baseUri, PIPE_NAME, authProvider);
+ webClientService = new StandardWebClientService();
+ client = new SnowpipeIngestClient(baseUri, PIPE_NAME, authProvider,
webClientService);
}
@AfterEach
void closeClient() {
- client.close();
+ webClientService.close();
}
@Test
diff --git
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/TestStandardSnowflakeIngestManagerProviderService.java
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/TestStandardSnowflakeIngestManagerProviderService.java
index cdb25ff4a02..c5defe1f11e 100644
---
a/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/TestStandardSnowflakeIngestManagerProviderService.java
+++
b/nifi-extension-bundles/nifi-snowflake-bundle/nifi-snowflake-services/src/test/java/org/apache/nifi/snowflake/service/TestStandardSnowflakeIngestManagerProviderService.java
@@ -28,6 +28,9 @@ import
org.apache.nifi.processors.snowflake.snowpipe.InsertReport;
import org.apache.nifi.processors.snowflake.util.SnowflakeProperties;
import org.apache.nifi.util.MockPropertyConfiguration;
import org.apache.nifi.util.PropertyMigrationResult;
+import org.apache.nifi.web.client.StandardWebClientService;
+import org.apache.nifi.web.client.api.WebClientService;
+import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
@@ -40,8 +43,10 @@ import java.security.NoSuchAlgorithmException;
import java.security.interfaces.RSAPrivateCrtKey;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -101,11 +106,19 @@ public class
TestStandardSnowflakeIngestManagerProviderService {
@StartStop
public final MockWebServer mockWebServer = new MockWebServer();
+ private StandardWebClientService webClientService;
+
private SnowpipeIngestClient client;
@BeforeEach
void setUp() throws NoSuchAlgorithmException {
- client = createClient(mockWebServer);
+ webClientService = new StandardWebClientService();
+ client = createClient(mockWebServer, webClientService);
+ }
+
+ @AfterEach
+ void closeClient() {
+ webClientService.close();
}
@Test
@@ -134,6 +147,12 @@ public class
TestStandardSnowflakeIngestManagerProviderService {
final Map<String, String> propertiesRenamed =
result.getPropertiesRenamed();
assertEquals(expectedRenamed, propertiesRenamed);
+
+ final Set<MockPropertyConfiguration.CreatedControllerService>
createdControllerServices = result.getCreatedControllerServices();
+ assertFalse(createdControllerServices.isEmpty());
+
+ final MockPropertyConfiguration.CreatedControllerService
createdControllerService = createdControllerServices.iterator().next();
+
assertTrue(createdControllerService.implementationClassName().endsWith("StandardWebClientServiceProvider"));
}
@Test
@@ -202,12 +221,12 @@ public class
TestStandardSnowflakeIngestManagerProviderService {
assertTrue(exception.getMessage().contains(String.valueOf(HttpURLConnection.HTTP_NOT_FOUND)));
}
- private static SnowpipeIngestClient createClient(final MockWebServer
mockWebServer) throws NoSuchAlgorithmException {
+ private static SnowpipeIngestClient createClient(final MockWebServer
mockWebServer, final WebClientService webClientService) throws
NoSuchAlgorithmException {
final URI baseUri =
URI.create(HTTP_URI_FORMAT.formatted(mockWebServer.getHostName(),
mockWebServer.getPort()));
final KeyPairGenerator keyPairGenerator =
KeyPairGenerator.getInstance(KEY_ALGORITHM);
final KeyPair keyPair = keyPairGenerator.generateKeyPair();
final RSAPrivateCrtKey privateKey = (RSAPrivateCrtKey)
keyPair.getPrivate();
final RSAKeyAuthorizationProvider authProvider = new
RSAKeyAuthorizationProvider(ACCOUNT, USER, privateKey);
- return new SnowpipeIngestClient(baseUri, PIPE_NAME, authProvider);
+ return new SnowpipeIngestClient(baseUri, PIPE_NAME, authProvider,
webClientService);
}
}
