diff --git a/build/Dockerfiles/dataproxy.Dockerfile b/build/Dockerfiles/dataproxy.Dockerfile index d5c26fa..eaec760 100644 --- a/build/Dockerfiles/dataproxy.Dockerfile +++ b/build/Dockerfiles/dataproxy.Dockerfile @@ -23,7 +23,13 @@ RUN ln -s ${JAVA_HOME}/lib/libjli.so /lib64 COPY dataproxy-server/target/dataproxy-server-0.0.1-SNAPSHOT.jar dataproxy.jar COPY libs/*.jar libs/ -ENV JAVA_OPTS="-server -XX:+UseG1GC -XX:+DisableExplicitGC -XX:InitiatingHeapOccupancyPercent=68 -Xlog:gc*=info:file=gc.log:time,tags:filecount=5,filesize=10M -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/app/log -XX:ErrorFile=/app/log/hs_err_pid%p.log" +# At present, when the concurrency of the task is not high, +# the heap memory consumption is not large, the maximum setting is 512M, +# if the resources are sufficient, and the concurrency can be appropriately increased, +# but the MaxDirectMemorySize needs to be increased according to the situation, +# the current default setting is 1536m, +# You can override this configuration by adding environment variables +ENV JAVA_OPTS="-server -XX:+UseG1GC -XX:+UseContainerSupport -Xms256m -Xmx512m -XX:MaxDirectMemorySize=1536m -XX:+DisableExplicitGC -XX:InitiatingHeapOccupancyPercent=68 -Xlog:gc*=info:file=gc.log:time,tags:filecount=5,filesize=10M -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/app/log -XX:ErrorFile=/app/log/hs_err_pid%p.log" ENV LOG_LEVEL=INFO EXPOSE 8023 ENTRYPOINT ${JAVA_HOME}/bin/java ${JAVA_OPTS} -Dsun.net.http.allowRestrictedHeaders=true --add-opens=java.base/java.nio=ALL-UNNAMED -jar ./dataproxy.jar \ No newline at end of file diff --git a/dataproxy-core/pom.xml b/dataproxy-core/pom.xml index 6b90222..bd91b48 100644 --- a/dataproxy-core/pom.xml +++ b/dataproxy-core/pom.xml @@ -33,10 +33,10 @@ org.apache.arrow arrow-format - + org.apache.arrow arrow-memory @@ -122,6 +122,12 @@ + + org.apache.arrow + arrow-memory-netty + test + + com.github.ben-manes.caffeine caffeine @@ -146,5 +152,28 @@ jsr305 + + + org.junit.jupiter + junit-jupiter-api + test + + + org.junit.jupiter + junit-jupiter-engine + test + + + org.mockito + mockito-inline + test + + + + + org.mockito + mockito-junit-jupiter + test + diff --git a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/DefaultFlightServerConfigLoader.java b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/DefaultFlightServerConfigLoader.java index bc32252..434e435 100644 --- a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/DefaultFlightServerConfigLoader.java +++ b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/DefaultFlightServerConfigLoader.java @@ -61,5 +61,6 @@ public void loadProperties(Properties properties) { throw new RuntimeException(e); } properties.put(FlightServerConfigKey.PORT, 8023); + properties.put(FlightServerConfigKey.METRICS_PORT, 9101); } } diff --git a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfig.java b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfig.java index 6e0a28f..2fa85ae 100644 --- a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfig.java +++ b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfig.java @@ -18,6 +18,9 @@ import lombok.extern.slf4j.Slf4j; import org.apache.arrow.flight.Location; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.secretflow.dataproxy.core.listener.DataProxyAllocationListener; /** * @author yuexie @@ -26,7 +29,14 @@ @Slf4j public record FlightServerConfig(String host, int port) { + private static final BufferAllocator ROOT_ALLOCATOR = + new RootAllocator(new DataProxyAllocationListener(), 2L * 1024 * 1024 * 1024); + public Location getLocation() { return Location.forGrpcInsecure(host, port); } + + public BufferAllocator getBufferAllocator() { + return ROOT_ALLOCATOR; + } } diff --git a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfigKey.java b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfigKey.java index 3ddd2cb..fae4732 100644 --- a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfigKey.java +++ b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/config/FlightServerConfigKey.java @@ -26,4 +26,6 @@ public class FlightServerConfigKey { public static final String PORT = "SERVICE_PORT"; + public static final String METRICS_PORT = "METRICS_PORT"; + } diff --git a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListener.java b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListener.java new file mode 100644 index 0000000..fb82966 --- /dev/null +++ b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListener.java @@ -0,0 +1,107 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.core.listener; + +import lombok.extern.slf4j.Slf4j; +import org.apache.arrow.memory.AllocationListener; +import org.apache.arrow.memory.AllocationOutcome; +import org.apache.arrow.memory.BufferAllocator; + +/** + * @author yuexie + * @date 2025/4/14 16:14 + **/ +@Slf4j +public class DataProxyAllocationListener implements AllocationListener { + + /** + * Called each time a new buffer has been requested. + * + *

An exception can be safely thrown by this method to terminate the allocation. + * + * @param size the buffer size being allocated + */ + @Override + public void onPreAllocation(long size) { + AllocationListener.super.onPreAllocation(size); + log.debug("onPreAllocation, size: {}", size); + } + + /** + * Called each time a new buffer has been allocated. + * + *

An exception cannot be thrown by this method. + * + * @param size the buffer size being allocated + */ + @Override + public void onAllocation(long size) { + AllocationListener.super.onAllocation(size); + log.debug("onAllocation, size: {}", size); + } + + /** + * Informed each time a buffer is released from allocation. + * + *

An exception cannot be thrown by this method. + * + * @param size The size of the buffer being released. + */ + @Override + public void onRelease(long size) { + AllocationListener.super.onRelease(size); + } + + /** + * Called whenever an allocation failed, giving the caller a chance to create some space in the + * allocator (either by freeing some resource, or by changing the limit), and, if successful, + * allowing the allocator to retry the allocation. + * + * @param size the buffer size that was being allocated + * @param outcome the outcome of the failed allocation. Carries information of what failed + * @return true, if the allocation can be retried; false if the allocation should fail + */ + @Override + public boolean onFailedAllocation(long size, AllocationOutcome outcome) { + log.debug("onFailedAllocation, size: {}, outcome: {}", size, outcome); + return AllocationListener.super.onFailedAllocation(size, outcome); + } + + /** + * Called immediately after a child allocator was added to the parent allocator. + * + * @param parentAllocator The parent allocator to which a child was added + * @param childAllocator The child allocator that was just added + */ + @Override + public void onChildAdded(BufferAllocator parentAllocator, BufferAllocator childAllocator) { + AllocationListener.super.onChildAdded(parentAllocator, childAllocator); + log.debug("onChildAdded, childAllocator: {}, size: {}", childAllocator.getName(), childAllocator.getLimit()); + } + + /** + * Called immediately after a child allocator was removed from the parent allocator. + * + * @param parentAllocator The parent allocator from which a child was removed + * @param childAllocator The child allocator that was just removed + */ + @Override + public void onChildRemoved(BufferAllocator parentAllocator, BufferAllocator childAllocator) { + AllocationListener.super.onChildRemoved(parentAllocator, childAllocator); + log.debug("onChildRemoved, childAllocator: {}, size: {}", childAllocator.getName(), childAllocator.getLimit()); + } +} diff --git a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/reader/AbstractSender.java b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/reader/AbstractSender.java index 400b4df..8b664b9 100644 --- a/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/reader/AbstractSender.java +++ b/dataproxy-core/src/main/java/org/secretflow/dataproxy/core/reader/AbstractSender.java @@ -81,8 +81,9 @@ public void send() { ValueVectorUtility.ensureCapacity(root, takeRecordCount + 1); this.toArrowVector(record, root, takeRecordCount); takeRecordCount++; - - if (takeRecordCount % 300_000 == 0) { + // 10w records, flush to arrow + // It can't be too big, and the off-heap memory is clipped + if (takeRecordCount % 100_000 == 0) { break; } } @@ -158,7 +159,7 @@ public void close() throws Exception { * Pre-application for arrow vector memory */ private void preAllocate() { - + root.clear(); ValueVectorUtility.preAllocate(root, estimatedRecordCount); root.getFieldVectors().forEach(fieldVector -> { @@ -168,6 +169,5 @@ private void preAllocate() { baseVariableWidthVector.allocateNew(estimatedRecordCount * 32); } }); - root.clear(); } } diff --git a/dataproxy-core/src/test/java/org/secretflow/dataproxy/core/config/FlightServerContextTest.java b/dataproxy-core/src/test/java/org/secretflow/dataproxy/core/config/FlightServerContextTest.java new file mode 100644 index 0000000..c98e31f --- /dev/null +++ b/dataproxy-core/src/test/java/org/secretflow/dataproxy/core/config/FlightServerContextTest.java @@ -0,0 +1,73 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.core.config; + +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +/** + * @author yuexie + * @date 2025/06/16 10:59:10 + */ +public class FlightServerContextTest { + + @Test + void testGetInstanceInMultiThread() throws InterruptedException { + final int threadCount = 10; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch latch = new CountDownLatch(threadCount); + AtomicReference firstInstance = new AtomicReference<>(); + + for (int i = 0; i < threadCount; i++) { + executor.submit(() -> { + FlightServerContext instance = FlightServerContext.getInstance(); + if (firstInstance.get() == null) { + firstInstance.set(instance); + } else { + assertSame(firstInstance.get(), instance, "The same instance should be returned in a multi-threaded environment"); + } + latch.countDown(); + }); + } + + latch.await(5, TimeUnit.SECONDS); + executor.shutdown(); + } + @Test + void testGetOrDefault() { + String testKey = "nonexistent.key"; + String defaultValue = "default"; + + String result = FlightServerContext.getOrDefault(testKey, String.class, defaultValue); + assertEquals(defaultValue, result); + } + + @Test + void testFlightServerConfig() { + FlightServerContext context = FlightServerContext.getInstance(); + assertNotNull(context.getFlightServerConfig(), "flightServerConfig should be initialized"); + } + +} diff --git a/dataproxy-core/src/test/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListenerTest.java b/dataproxy-core/src/test/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListenerTest.java new file mode 100644 index 0000000..494b80d --- /dev/null +++ b/dataproxy-core/src/test/java/org/secretflow/dataproxy/core/listener/DataProxyAllocationListenerTest.java @@ -0,0 +1,82 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.core.listener; + +import org.apache.arrow.memory.AllocationOutcome; +import org.apache.arrow.memory.BufferAllocator; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertFalse; + +public class DataProxyAllocationListenerTest { + + private final DataProxyAllocationListener listener = new DataProxyAllocationListener(); + + @Test + public void testOnPreAllocation() { + assertDoesNotThrow(() -> listener.onPreAllocation(1024L)); + } + + @Test + public void testOnAllocation() { + assertDoesNotThrow(() -> listener.onAllocation(1024L)); + } + + @Test + public void testOnRelease() { + assertDoesNotThrow(() -> listener.onRelease(1024L)); + } + + @Test + public void testOnFailedAllocation_returnsFalseByDefault() { + AllocationOutcome outcome = Mockito.mock(AllocationOutcome.class); + assertDoesNotThrow(() -> { + boolean result = listener.onFailedAllocation(1024L, outcome); + assertFalse(result); + }); + } + + @Test + public void testOnChildAdded() { + try (BufferAllocator parentAllocator = Mockito.mock(BufferAllocator.class); + BufferAllocator childAllocator = Mockito.mock(BufferAllocator.class)) { + + Mockito.when(parentAllocator.getName()).thenReturn("parentAllocator"); + Mockito.when(childAllocator.getName()).thenReturn("childAllocator"); + Mockito.when(childAllocator.getLimit()).thenReturn(1024L); + + assertDoesNotThrow(() -> listener.onChildAdded(parentAllocator, childAllocator)); + } + } + + @Test + public void testOnChildRemoved() { + + try (BufferAllocator parentAllocator = Mockito.mock(BufferAllocator.class); + BufferAllocator childAllocator = Mockito.mock(BufferAllocator.class)) { + + Mockito.when(parentAllocator.getName()).thenReturn("parentAllocator"); + Mockito.when(childAllocator.getName()).thenReturn("childAllocator"); + Mockito.when(childAllocator.getLimit()).thenReturn(1024L); + + assertDoesNotThrow(() -> listener.onChildRemoved(parentAllocator, childAllocator)); + } + } + + +} \ No newline at end of file diff --git a/dataproxy-integration-tests/src/test/java/org/secretflow/dataproxy/integration/tests/OdpsIntegrationTest.java b/dataproxy-integration-tests/src/test/java/org/secretflow/dataproxy/integration/tests/OdpsIntegrationTest.java index 1bae754..fa85117 100644 --- a/dataproxy-integration-tests/src/test/java/org/secretflow/dataproxy/integration/tests/OdpsIntegrationTest.java +++ b/dataproxy-integration-tests/src/test/java/org/secretflow/dataproxy/integration/tests/OdpsIntegrationTest.java @@ -43,6 +43,8 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.MethodOrderer; import org.junit.jupiter.api.Order; import org.junit.jupiter.api.Test; @@ -50,7 +52,9 @@ import org.junit.jupiter.api.condition.EnabledIfSystemProperty; import org.junit.jupiter.api.io.TempDir; import org.secretflow.dataproxy.common.utils.ArrowUtil; +import org.secretflow.dataproxy.core.config.FlightServerContext; import org.secretflow.dataproxy.integration.tests.utils.OdpsTestUtil; +import org.secretflow.dataproxy.server.DataProxyFlightServer; import org.secretflow.v1alpha1.common.Common; import org.secretflow.v1alpha1.kusciaapi.Domaindata; import org.secretflow.v1alpha1.kusciaapi.Domaindatasource; @@ -81,8 +85,10 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; /** * @author yuexie @@ -147,6 +153,44 @@ public class OdpsIntegrationTest extends BaseArrowFlightServerTest { private static Path tempDir; private static Path tmpFilePath; + @BeforeAll + public static void startServer() { + + assertNotEquals("", OdpsTestUtil.getOdpsProject(), "odps project is empty"); + assertNotEquals("", OdpsTestUtil.getOdpsEndpoint(), "odps endpoint is empty"); + assertNotEquals("", OdpsTestUtil.getAccessKeyId(), "odps access key id is empty"); + assertNotEquals("", OdpsTestUtil.getAccessKeySecret(), "odps access key secret is empty"); + + dataProxyFlightServer = new DataProxyFlightServer(FlightServerContext.getInstance().getFlightServerConfig()); + + assertDoesNotThrow(() -> { + serverThread = new Thread(() -> { + try { + dataProxyFlightServer.start(); + SERVER_START_LATCH.countDown(); + dataProxyFlightServer.awaitTermination(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Exception e) { + fail("Exception was thrown: " + e.getMessage()); + } + }); + }); + + assertDoesNotThrow(() -> { + serverThread.start(); + SERVER_START_LATCH.await(); + }); + } + + @AfterAll + static void stopServer() { + assertDoesNotThrow(() -> { + if (dataProxyFlightServer != null) dataProxyFlightServer.close(); + serverThread.interrupt(); + }); + } + @Test @Order(2) public void testDoGetWithTable() { @@ -466,4 +510,4 @@ private static String bytesToHex(byte[] bytes) { } return hexString.toString(); } -} \ No newline at end of file +} diff --git a/dataproxy-metrics/pom.xml b/dataproxy-metrics/pom.xml new file mode 100644 index 0000000..b476221 --- /dev/null +++ b/dataproxy-metrics/pom.xml @@ -0,0 +1,67 @@ + + + + + 4.0.0 + + org.secretflow + dataproxy + 0.0.1-SNAPSHOT + + + dataproxy-metrics + + + + + io.micrometer + micrometer-core + + + io.micrometer + micrometer-registry-prometheus + + + + + org.junit.jupiter + junit-jupiter-api + test + + + org.junit.jupiter + junit-jupiter-engine + test + + + org.mockito + mockito-inline + test + + + + + org.mockito + mockito-junit-jupiter + test + + + + + \ No newline at end of file diff --git a/dataproxy-metrics/src/main/java/org/secretflow/dataproxy/metrics/JvmMetricsRegistrar.java b/dataproxy-metrics/src/main/java/org/secretflow/dataproxy/metrics/JvmMetricsRegistrar.java new file mode 100644 index 0000000..b86c5dc --- /dev/null +++ b/dataproxy-metrics/src/main/java/org/secretflow/dataproxy/metrics/JvmMetricsRegistrar.java @@ -0,0 +1,50 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.metrics; + +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.binder.jvm.ClassLoaderMetrics; +import io.micrometer.core.instrument.binder.jvm.JvmGcMetrics; +import io.micrometer.core.instrument.binder.jvm.JvmMemoryMetrics; +import io.micrometer.core.instrument.binder.jvm.JvmThreadMetrics; +import io.micrometer.core.instrument.binder.system.FileDescriptorMetrics; + +/** + * @author yuexie + * @date 2025/4/14 10:46 + **/ +public class JvmMetricsRegistrar { + public static void registerJvmMemoryMetrics(MeterRegistry registry) { + new JvmMemoryMetrics().bindTo(registry); + } + + public static void registerJvmThreadMetrics(MeterRegistry registry) { + new JvmThreadMetrics().bindTo(registry); + } + + public static void registerJvmGcMetrics(MeterRegistry registry) { + try (JvmGcMetrics jvmGcMetrics = new JvmGcMetrics()) { + jvmGcMetrics.bindTo(registry); + } + } + public static void registerClassLoaderMetrics(MeterRegistry registry) { + new ClassLoaderMetrics().bindTo(registry); + } + public static void registerFileDescriptorMetrics(MeterRegistry registry) { + new FileDescriptorMetrics().bindTo(registry); + } +} diff --git a/dataproxy-metrics/src/test/java/org/secretflow/dataproxy/metrics/JvmMetricsRegistrarTest.java b/dataproxy-metrics/src/test/java/org/secretflow/dataproxy/metrics/JvmMetricsRegistrarTest.java new file mode 100644 index 0000000..28d34d0 --- /dev/null +++ b/dataproxy-metrics/src/test/java/org/secretflow/dataproxy/metrics/JvmMetricsRegistrarTest.java @@ -0,0 +1,140 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.metrics; + +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.binder.jvm.ClassLoaderMetrics; +import io.micrometer.core.instrument.binder.jvm.JvmGcMetrics; +import io.micrometer.core.instrument.binder.jvm.JvmMemoryMetrics; +import io.micrometer.core.instrument.binder.jvm.JvmThreadMetrics; +import io.micrometer.core.instrument.binder.system.FileDescriptorMetrics; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.MockedConstruction; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.verify; + +/** + * @author yuexie + * @date 2025/6/3 11:16 + **/ +@ExtendWith(MockitoExtension.class) +public class JvmMetricsRegistrarTest { + + /** + * Test Scenario: Under normal circumstances, the JVM memory metric is registered + */ + @Test + public void testRegisterJvmMemoryMetrics_success() { + MeterRegistry registry = Mockito.mock(MeterRegistry.class); + try (MockedConstruction mocked = Mockito.mockConstruction(JvmMemoryMetrics.class)) { + JvmMetricsRegistrar.registerJvmMemoryMetrics(registry); + verify(mocked.constructed().get(0)).bindTo(registry); + } + } + + /** + * Test Scenario: MeterRegistry is null, which registers JVM memory metrics + */ + @Test + public void testRegisterJvmMemoryMetrics_nullRegistry() { + assertThrows(NullPointerException.class, () -> JvmMetricsRegistrar.registerJvmMemoryMetrics(null)); + } + + /** + * Test Scenario: Normal, register the JVM thread metric + */ + @Test + public void testRegisterJvmThreadMetrics_success() { + MeterRegistry registry = Mockito.mock(MeterRegistry.class); + try (MockedConstruction mocked = Mockito.mockConstruction(JvmThreadMetrics.class)) { + JvmMetricsRegistrar.registerJvmThreadMetrics(registry); + verify(mocked.constructed().get(0)).bindTo(registry); + } + } + + /** + * Test Scenario: MeterRegistry is null, which registers JVM thread metrics + */ + @Test + public void testRegisterJvmThreadMetrics_nullRegistry() { + assertThrows(NullPointerException.class, () -> JvmMetricsRegistrar.registerJvmThreadMetrics(null)); + } + + /** + * Test Scenario: Normal, register the JVM garbage collection metric + */ + @Test + public void testRegisterJvmGcMetrics_success() { + MeterRegistry registry = Mockito.mock(MeterRegistry.class); + try (MockedConstruction mocked = Mockito.mockConstruction(JvmGcMetrics.class)) { + JvmMetricsRegistrar.registerJvmGcMetrics(registry); + verify(mocked.constructed().get(0)).bindTo(registry); + } + } + + /** + * Test scenario: MeterRegistry is null and the JVM garbage collection metric is registered + */ + @Test + public void testRegisterJvmGcMetrics_nullRegistry() { + assertThrows(NullPointerException.class, () -> JvmMetricsRegistrar.registerJvmGcMetrics(null)); + } + + /** + * Test scenario: Normal, register the classloader metric + */ + @Test + public void testRegisterClassLoaderMetrics_success() { + MeterRegistry registry = Mockito.mock(MeterRegistry.class); + try (MockedConstruction mocked = Mockito.mockConstruction(ClassLoaderMetrics.class)) { + JvmMetricsRegistrar.registerClassLoaderMetrics(registry); + verify(mocked.constructed().get(0)).bindTo(registry); + } + } + + /** + * Test scenario: MeterRegistry is null, and the classloader indicator is registered + */ + @Test + public void testRegisterClassLoaderMetrics_nullRegistry() { + assertThrows(NullPointerException.class, () -> JvmMetricsRegistrar.registerClassLoaderMetrics(null)); + } + + /** + * Test Scenario: Normal, register file descriptor metrics + */ + @Test + public void testRegisterFileDescriptorMetrics_success() { + MeterRegistry registry = Mockito.mock(MeterRegistry.class); + try (MockedConstruction mocked = Mockito.mockConstruction(FileDescriptorMetrics.class)) { + JvmMetricsRegistrar.registerFileDescriptorMetrics(registry); + verify(mocked.constructed().get(0)).bindTo(registry); + } + } + + /** + * Test scenario: MeterRegistry is null, and the file descriptor indicator is registered + */ + @Test + public void testRegisterFileDescriptorMetrics_nullRegistry() { + assertThrows(NullPointerException.class, () -> JvmMetricsRegistrar.registerFileDescriptorMetrics(null)); + } +} diff --git a/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/producer/AbstractDatabaseFlightProducer.java b/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/producer/AbstractDatabaseFlightProducer.java index 466f6f7..c71d172 100644 --- a/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/producer/AbstractDatabaseFlightProducer.java +++ b/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/producer/AbstractDatabaseFlightProducer.java @@ -141,6 +141,7 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l break; } } + log.info("doGet is completed"); listener.completed(); } catch (InterruptedException | IOException e) { log.error("doGet error!", e); @@ -167,6 +168,7 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream ) { + log.info("acceptPut hive write."); final Any any = GrpcUtils.parseOrThrow(flightStream.getDescriptor().getCommand()); if(!"type.googleapis.com/kuscia.proto.api.v1alpha1.datamesh.TicketDomainDataQuery".equals(any.getTypeUrl())) { @@ -203,6 +205,12 @@ public Runnable acceptPut( .withCause(e) .withDescription(e.getMessage()) .toRuntimeException(); + } catch (Exception e) { + log.error("unknown error", e); + throw CallStatus.INTERNAL + .withCause(e) + .withDescription(e.getMessage()) + .toRuntimeException(); } }; } diff --git a/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/reader/DatabaseRecordReader.java b/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/reader/DatabaseRecordReader.java index 3ad1bd6..e64fdf7 100644 --- a/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/reader/DatabaseRecordReader.java +++ b/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/reader/DatabaseRecordReader.java @@ -23,7 +23,6 @@ import org.secretflow.dataproxy.core.reader.Sender; import java.sql.ResultSet; -import java.sql.SQLException; @Slf4j public class DatabaseRecordReader extends AbstractReader { diff --git a/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/writer/DatabaseRecordWriter.java b/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/writer/DatabaseRecordWriter.java index 2f5b8b2..5976a1a 100644 --- a/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/writer/DatabaseRecordWriter.java +++ b/dataproxy-plugins/dataproxy-plugin-database/src/main/java/org/secretflow/dataproxy/plugin/database/writer/DatabaseRecordWriter.java @@ -45,6 +45,7 @@ public class DatabaseRecordWriter implements Writer { private final Function initFunc; private final BiFunction checkTableExists; private final int BATCH_NUM = 500; + @FunctionalInterface public interface BuildCreateTableSqlFunc { String apply(String tableName, Schema schema, Map partitionSpec); @@ -58,34 +59,18 @@ public interface BuildInsertSqlFunc { @FunctionalInterface public interface BuildMultiInsertSqlFunc { - String apply(String tableName, Schema schema, List> data, Map partitionSpec); + SqlWithParams apply(String tableName, Schema schema, List> data, Map partitionSpec); } private BuildMultiInsertSqlFunc buildMultiInsertSql; private final Map partitionSpec; private final String tableName; private Connection connection; - private final boolean supportMutliInsert; - - public static Map PasrsePartition(String partition) { - Map res = new LinkedHashMap<>(); - String[] groups = partition.split("[,/]"); - for (String group : groups) { - String[] kv = group.split("="); - if (kv.length != 2) { - throw new IllegalArgumentException("Invalid partition spec."); - } - - String k = kv[0].trim(); - String v = kv[1].trim() - .replaceAll("'", "") - .replaceAll("\"", ""); - if (k.isEmpty() || v.isEmpty()) { - throw new IllegalArgumentException("Invalid partition spec."); - } + private final boolean supportMultiInsert; - res.put(k, v); - } - return res; + // TODO: parse partition + // Partitioning is not supported in the current version + public static Map parsePartition(String partition) { + return new LinkedHashMap<>(); } public DatabaseRecordWriter(DatabaseWriteConfig commandConfig, @@ -101,8 +86,8 @@ public DatabaseRecordWriter(DatabaseWriteConfig commandConfig, this.buildCreateTableSql = buildCreateTableSql; this.buildInsertSql = buildInsertSql; this.tableName = this.dbTableConfig.tableName(); - this.partitionSpec = PasrsePartition(this.dbTableConfig.partition()); - supportMutliInsert = false; + this.partitionSpec = parsePartition(this.dbTableConfig.partition()); + supportMultiInsert = false; this.prepare(); } @@ -119,8 +104,8 @@ public DatabaseRecordWriter(DatabaseWriteConfig commandConfig, this.buildCreateTableSql = buildCreateTableSql; this.buildMultiInsertSql = buildMultiInsertSql; this.tableName = this.dbTableConfig.tableName(); - this.partitionSpec = PasrsePartition(this.dbTableConfig.partition()); - supportMutliInsert = true; + this.partitionSpec = parsePartition(this.dbTableConfig.partition()); + supportMultiInsert = true; this.prepare(); } @@ -140,7 +125,7 @@ private void prepare(){ } /** - * 获取字段数据 + * Get field data * * @param fieldVector field vector * @param index index @@ -196,7 +181,7 @@ public void write(VectorSchemaRoot root) { String columnName; - if(supportMutliInsert) { + if(supportMultiInsert) { List> multiRecords = new ArrayList<>(); for(int rowIndex = 0; rowIndex < batchSize; rowIndex ++) { Record record = new Record(); @@ -249,10 +234,18 @@ private void createTable(Schema schema){ } - private void dropTable() throws SQLException { + private void validateTableName(String tableName) { if (tableName == null || tableName.trim().isEmpty()) { throw new IllegalArgumentException("Table name cannot be null or empty"); } + // Only allows letters, numbers, underscores, and must start with a letter + if (!tableName.matches("^[a-zA-Z][a-zA-Z0-9_]*$")) { + throw new IllegalArgumentException("Invalid table name format. Table name must start with a letter and contain only letters, numbers, and underscores"); + } + } + + private void dropTable() throws SQLException { + validateTableName(tableName); String sql = "DROP TABLE IF EXISTS " + tableName; @@ -266,9 +259,7 @@ private void dropTable() throws SQLException { } private void deleteAllRowOfTable() throws SQLException { - if (tableName == null || tableName.trim().isEmpty()) { - throw new IllegalArgumentException("Table name cannot be null or empty"); - } + validateTableName(tableName); String sql = "DELETE FROM " + tableName; @@ -292,11 +283,15 @@ public void insertData(Schema arrowSchema, Map data) { } public void insertMultiData(Schema arrowSchema, List> multiData){ - String sql = this.buildMultiInsertSql.apply(tableName, arrowSchema, multiData, partitionSpec); - try (PreparedStatement stmt = connection.prepareStatement(sql)){ - stmt.executeUpdate(); + + SqlWithParams sp = this.buildMultiInsertSql.apply(tableName, arrowSchema, multiData, partitionSpec); + try (PreparedStatement ps = connection.prepareStatement(sp.sql);){ + for (int i = 0; i < sp.params.size(); i++) { + ps.setObject(i + 1, sp.params.get(i)); + } + ps.executeUpdate(); } catch (SQLException e) { - log.error("insert data error: sql:\"{}\" error:\"{}\"", sql, e.getMessage()); + log.error("insert data error: sql:\"{}\" error:\"{}\"", sp, e.getMessage()); throw new RuntimeException(e); } } @@ -323,4 +318,15 @@ private void preProcessing(String tableName){ createTable(commandConfig.getResultSchema()); } + public static class SqlWithParams { + + public final String sql; + public final List params; + + public SqlWithParams(String sql, List params) { + this.sql = sql; + this.params = params; + } + } + } diff --git a/dataproxy-plugins/dataproxy-plugin-hive/pom.xml b/dataproxy-plugins/dataproxy-plugin-hive/pom.xml index 80f7e58..576197f 100644 --- a/dataproxy-plugins/dataproxy-plugin-hive/pom.xml +++ b/dataproxy-plugins/dataproxy-plugin-hive/pom.xml @@ -108,6 +108,11 @@ lombok compile + + org.projectlombok + lombok + compile + diff --git a/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducer.java b/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducer.java index 5299054..6c13653 100644 --- a/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducer.java +++ b/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducer.java @@ -37,11 +37,13 @@ protected DatabaseDoGetContext initDoGetContext(DatabaseCommandConfig config) @Override protected DatabaseRecordWriter initRecordWriter(DatabaseWriteConfig config) { + return new DatabaseRecordWriter(config, HiveUtil::initHive, HiveUtil::buildCreateTableSql, HiveUtil::buildMultiRowInsertSql, - HiveUtil::checkTableExists); + HiveUtil::checkTableExists + ); } } diff --git a/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtil.java b/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtil.java index 526c0cb..dfb48df 100644 --- a/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtil.java +++ b/dataproxy-plugins/dataproxy-plugin-hive/src/main/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtil.java @@ -28,17 +28,26 @@ import org.secretflow.dataproxy.plugin.database.config.DatabaseConnectConfig; import org.secretflow.dataproxy.plugin.database.writer.DatabaseRecordWriter; -import java.sql.*; -import java.util.*; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; -import static org.secretflow.dataproxy.plugin.database.writer.DatabaseRecordWriter.PasrsePartition; +import static org.secretflow.dataproxy.plugin.database.writer.DatabaseRecordWriter.parsePartition; @Slf4j public class HiveUtil { - + public static Connection initHive(DatabaseConnectConfig config) { String endpoint = config.endpoint(); String ip; @@ -53,13 +62,27 @@ public static Connection initHive(DatabaseConnectConfig config) { } else { ip = endpoint; } + + // Validate IP address/hostname to prevent JDBC URL injection + if (ip == null || !ip.matches("^[a-zA-Z0-9._-]+$") || ip.contains("..") || ip.startsWith(".") || ip.endsWith(".")) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, + "Invalid IP address or hostname: " + ip); + } + + // Validate database name to prevent JDBC URL injection + String database = config.database(); + if (database == null || !database.matches("^[a-zA-Z0-9_]+$")) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, + "Invalid database name: " + database); + } + Connection conn; try{ // hive Authentication None if(!config.username().isEmpty() && !config.password().isEmpty()) { - conn = DriverManager.getConnection(String.format("jdbc:hive2://%s:%s/%s", ip, port, config.database()), config.username(), config.password()); + conn = DriverManager.getConnection(String.format("jdbc:hive2://%s:%s/%s", ip, port, database), config.username(), config.password()); } else { - conn = DriverManager.getConnection(String.format("jdbc:hive2://%s:%s/%s", ip, port, config.database())); + conn = DriverManager.getConnection(String.format("jdbc:hive2://%s:%s/%s", ip, port, database)); } } catch (Exception e) { log.error("database init error \"{}\"", e.getMessage()); @@ -70,53 +93,100 @@ public static Connection initHive(DatabaseConnectConfig config) { } public static String buildQuerySql(String tableName, List fields, String whereClause) { - final Pattern columnOrValuePattern = Pattern.compile("^[\\u00b7A-Za-z0-9\\u4e00-\\u9fa5\\-_,.]*$"); + final Pattern columnOrValuePattern = Pattern.compile("^[a-zA-Z0-9_]+$"); + // Validate table name if (!columnOrValuePattern.matcher(tableName).matches()) { throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid tableName:" + tableName); } + + // Validate field names + for (String field : fields) { + if (!columnOrValuePattern.matcher(field).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid field name:" + field); + } + } + + // Process where clause + String processedWhereClause = ""; if (!whereClause.isEmpty()) { String[] groups = whereClause.split("[,/]"); if (groups.length > 1) { - final Map partitionSpec = PasrsePartition(whereClause); + final Map partitionSpec = parsePartition(whereClause); - for (String key : partitionSpec.keySet()) { + for (Map.Entry entry : partitionSpec.entrySet()) { + String key = entry.getKey(); + String value = entry.getValue(); + + // Validate partition key name if (!columnOrValuePattern.matcher(key).matches()) { throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid partition key:" + key); } - if (!columnOrValuePattern.matcher(partitionSpec.get(key)).matches()) { - throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid partition value:" + partitionSpec.get(key)); + + // Validate partition value - only allow letters, numbers, underscores, hyphens, and dots + if (!value.matches("^[a-zA-Z0-9_.-]+$")) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid partition value:" + value); } } - List list = partitionSpec.keySet().stream().map(k -> k + "='" + partitionSpec.get(k) + "'").toList(); - whereClause = String.join(" and ", list); + List list = partitionSpec.keySet().stream() + .map(k -> k + "='" + escapeString(partitionSpec.get(k)) + "'") + .toList(); + processedWhereClause = String.join(" and ", list); + } else { + // For simple where conditions, use stricter validation or disallow + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, + "Invalid where clause format. Use partition format like 'key=value,key2=value2'"); } } - String sql = "select " + String.join(",", fields) + " from " + tableName + (whereClause.isEmpty() ? "" : " where " + whereClause); + + String sql = "select " + String.join(",", fields) + " from " + tableName + + (processedWhereClause.isEmpty() ? "" : " where " + processedWhereClause); log.info("buildQuerySql sql:{}", sql); return sql; } public static String buildCreateTableSql(String tableName, Schema schema, Map partition) { + final Pattern identifierPattern = Pattern.compile("^[a-zA-Z0-9_]+$"); + + // Validate table name + if (!identifierPattern.matcher(tableName).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid tableName:" + tableName); + } + StringBuilder sb = new StringBuilder(); sb.append("CREATE TABLE ").append(tableName).append(" (\n"); List fields = schema.getFields(); - Set partitionKeys = partition.keySet(); // 分区字段名集合 + Set partitionKeys = partition.keySet(); // Partition field names collection - // 用于快速通过字段名查找 Field + // Used for quick field lookup by name Map fieldMap = new LinkedHashMap<>(); for (Field field : fields) { fieldMap.put(field.getName(), field); } - // 表字段（不包含分区字段） + // Validate all field names + for (Field field : fields) { + String fieldName = field.getName(); + if (!identifierPattern.matcher(fieldName).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid field name:" + fieldName); + } + } + + // Validate partition key names + for (String partKey : partitionKeys) { + if (!identifierPattern.matcher(partKey).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid partition key:" + partKey); + } + } + + // Table fields (excluding partition fields) boolean first = true; for (Field field : fields) { String fieldName = field.getName(); if (partitionKeys.contains(fieldName)) { - continue; // 跳过分区字段 + continue; // Skip partition fields } if (!first) { @@ -129,14 +199,14 @@ public static String buildCreateTableSql(String tableName, Schema schema, Map data, Map partition) { - List fields = schema.getFields(); - Set partitionKeys = partition.keySet(); + public static DatabaseRecordWriter.SqlWithParams buildMultiRowInsertSql(String tableName, + Schema schema, + List> dataList, + Map partition + ) { + final Pattern identifierPattern = Pattern.compile("^[a-zA-Z0-9_]+$"); - List columns = new ArrayList<>(); - List values = new ArrayList<>(); - - // 遍历 schema 中的字段（以保持字段顺序） - for (Field field : fields) { - String fieldName = field.getName(); + if (dataList == null || dataList.isEmpty()) { + throw new IllegalArgumentException("No data to insert"); + } - // 分区字段单独处理 - if (partitionKeys.contains(fieldName)) { - continue; + if (!identifierPattern.matcher(tableName).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid tableName:" + tableName); + } + for (Field f : schema.getFields()) { + if (!identifierPattern.matcher(f.getName()).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid field name:" + f.getName()); } - - columns.add(fieldName); - Object rawValue = data.get(fieldName); - ArrowType arrowType = field.getType(); - - values.add(formatValue(rawValue, arrowType)); } - - // 构建 PARTITION 字段 - List partitionClauses = new ArrayList<>(); - for (String partKey : partitionKeys) { - Object partVal = data.get(partKey); // 注意：从 data 中获取值更安全 - Field field = fields.stream() - .filter(f -> f.getName().equals(partKey)) - .findFirst() - .orElseThrow(() -> new IllegalArgumentException("Partition key not found in schema: " + partKey)); - - String formatted = formatValue(partVal, field.getType()); - partitionClauses.add(partKey + "=" + formatted); + for (String k : partition.keySet()) { + if (!identifierPattern.matcher(k).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid partition key:" + k); + } } + Set partitionKeys = partition.keySet(); + List columns = schema.getFields().stream() + .map(Field::getName) + .filter(n -> !partitionKeys.contains(n)) + .collect(Collectors.toList()); + StringBuilder sb = new StringBuilder(); sb.append("INSERT INTO TABLE ").append(tableName); - if (!partitionClauses.isEmpty()) { - sb.append(" PARTITION (").append(String.join(", ", partitionClauses)).append(")"); - } - sb.append(" (").append(String.join(", ", columns)).append(")"); - sb.append(" VALUES (").append(String.join(", ", values)).append(")"); - log.info("buildInsertSql sql: {}", sb); - return sb.toString(); - } - - public static String buildMultiRowInsertSql(String tableName, Schema schema, List> dataList, Map partition) { - if (dataList == null || dataList.isEmpty()) { - throw new IllegalArgumentException("No data to insert"); + if (!partition.isEmpty()) { + List partitionPlaceholders = partition.keySet() + .stream() + .map(k -> k + " = ?") + .collect(Collectors.toList()); + sb.append(" PARTITION (").append(String.join(", ", partitionPlaceholders)).append(")"); } - Set partitionKeys = partition.keySet(); - List fields = schema.getFields(); + List colPlaceholders = columns.stream() + .map(c -> "?") + .collect(Collectors.toList()); + sb.append(" (").append(String.join(", ", columns)).append(") VALUES "); - // 取第一条数据判断 partition - Map firstRow = dataList.get(0); - List partitionClauses = new ArrayList<>(); - for (String partKey : partitionKeys) { - Object partVal = firstRow.get(partKey); - Field field = fields.stream() - .filter(f -> f.getName().equals(partKey)) - .findFirst() - .orElseThrow(() -> new IllegalArgumentException("Partition key not found in schema: " + partKey)); - partitionClauses.add(partKey + "=" + formatValue(partVal, field.getType())); - } + String singleRow = "(" + String.join(", ", colPlaceholders) + ")"; + List allRows = Collections.nCopies(dataList.size(), singleRow); + sb.append(String.join(", ", allRows)); - // 非 partition 字段 - List columns = fields.stream() - .map(Field::getName) - .filter(name -> !partitionKeys.contains(name)) - .collect(Collectors.toList()); + List params = new ArrayList<>(); - StringBuilder sb = new StringBuilder(); - sb.append("INSERT INTO TABLE ").append(tableName); - if (!partitionClauses.isEmpty()) { - sb.append(" PARTITION (").append(String.join(", ", partitionClauses)).append(")"); + for (String k : partition.keySet()) { + params.add(partition.get(k)); } - sb.append(" (").append(String.join(", ", columns)).append(")"); - sb.append(" VALUES "); - List rows = new ArrayList<>(); for (Map row : dataList) { - List vals = new ArrayList<>(); for (String col : columns) { - Object val = row.get(col); - ArrowType type = schema.findField(col).getType(); - vals.add(formatValue(val, type)); + params.add(row.get(col)); } - rows.add("(" + String.join(", ", vals) + ")"); - } - - sb.append(String.join(",\n", rows)); - return sb.toString(); - } - - private static String formatValue(Object value, ArrowType type) { - if (value == null) { - return "NULL"; } - return switch (type.getTypeID()) { - case Utf8, Binary, FixedSizeBinary -> "'" + escapeString(value.toString()) + "'"; - case Int, FloatingPoint, Bool -> value.toString(); - case Date, Timestamp, Time -> "'" + value + "'"; - case Decimal -> value.toString(); - default -> "'" + escapeString(value.toString()) + "'"; - }; + return new DatabaseRecordWriter.SqlWithParams(sb.toString(), params); } private static String escapeString(String str) { @@ -277,7 +302,7 @@ public static ArrowType jdbcType2ArrowType(String jdbcType) { String type = jdbcType.trim().toLowerCase(); - // 处理 decimal(p,s) + // Handle decimal(p,s) if (type.startsWith("decimal")) { Pattern pattern = Pattern.compile("decimal\\((\\d+),(\\d+)\\)"); Matcher matcher = pattern.matcher(type); @@ -286,7 +311,7 @@ public static ArrowType jdbcType2ArrowType(String jdbcType) { int scale = Integer.parseInt(matcher.group(2)); return new ArrowType.Decimal(precision, scale, 128); } else { - return new ArrowType.Decimal(38, 10, 128); // 默认精度 + return new ArrowType.Decimal(38, 10, 128); // Default precision } } @@ -310,32 +335,22 @@ public static ArrowType jdbcType2ArrowType(String jdbcType) { public static String arrowTypeToJdbcType(ArrowType arrowType) { if (arrowType instanceof ArrowType.Utf8) { return "STRING"; - } else if (arrowType instanceof ArrowType.Int) { - ArrowType.Int intType = (ArrowType.Int) arrowType; + } else if (arrowType instanceof ArrowType.Int intType) { int bitWidth = intType.getBitWidth(); boolean signed = intType.getIsSigned(); - switch (bitWidth) { - case 8: - return signed ? "TINYINT" : "TINYINT UNSIGNED"; - case 16: - return signed ? "SMALLINT" : "SMALLINT UNSIGNED"; - case 32: - return signed ? "INT" : "INT UNSIGNED"; - case 64: - return signed ? "BIGINT" : "BIGINT UNSIGNED"; - default: - throw new IllegalArgumentException("Unsupported Int bitWidth: " + bitWidth); - } - } else if (arrowType instanceof ArrowType.FloatingPoint) { - ArrowType.FloatingPoint fp = (ArrowType.FloatingPoint) arrowType; - switch (fp.getPrecision()) { - case SINGLE: - return "FLOAT"; - case DOUBLE: - return "DOUBLE"; - default: - throw new IllegalArgumentException("Unsupported floating point type"); - } + return switch (bitWidth) { + case 8 -> signed ? "TINYINT" : "TINYINT UNSIGNED"; + case 16 -> signed ? "SMALLINT" : "SMALLINT UNSIGNED"; + case 32 -> signed ? "INT" : "INT UNSIGNED"; + case 64 -> signed ? "BIGINT" : "BIGINT UNSIGNED"; + default -> throw new IllegalArgumentException("Unsupported Int bitWidth: " + bitWidth); + }; + } else if (arrowType instanceof ArrowType.FloatingPoint fp) { + return switch (fp.getPrecision()) { + case SINGLE -> "FLOAT"; + case DOUBLE -> "DOUBLE"; + default -> throw new IllegalArgumentException("Unsupported floating point type"); + }; } else if (arrowType instanceof ArrowType.Bool) { return "BOOLEAN"; } else if (arrowType instanceof ArrowType.Date) { @@ -344,8 +359,7 @@ public static String arrowTypeToJdbcType(ArrowType arrowType) { return "TIME"; } else if (arrowType instanceof ArrowType.Timestamp) { return "TIMESTAMP"; - } else if (arrowType instanceof ArrowType.Decimal) { - ArrowType.Decimal dec = (ArrowType.Decimal) arrowType; + } else if (arrowType instanceof ArrowType.Decimal dec) { return "DECIMAL(" + dec.getPrecision() + ", " + dec.getScale() + ")"; } else if (arrowType instanceof ArrowType.Binary || arrowType instanceof ArrowType.FixedSizeBinary) { return "BINARY"; @@ -355,16 +369,20 @@ public static String arrowTypeToJdbcType(ArrowType arrowType) { } public static boolean checkTableExists(Connection connection, String tableName) { + final Pattern identifierPattern = Pattern.compile("^[a-zA-Z0-9_]+$"); + + // Validate table name + if (!identifierPattern.matcher(tableName).matches()) { + throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "Invalid tableName:" + tableName); + } + ResultSet rs = null; - PreparedStatement stmt = null; + Statement stmt = null; try { - stmt = connection.prepareStatement("SHOW TABLES ?"); - stmt.setString(1, tableName); - rs = stmt.executeQuery(); - boolean exists = rs.next(); - rs.close(); - stmt.close(); - return exists; + // Hive doesn't support parameterized SHOW TABLES query, use string concatenation but validate table name first + stmt = connection.createStatement(); + rs = stmt.executeQuery("SHOW TABLES LIKE '" + tableName + "'"); + return rs.next(); } catch (SQLException e) { log.error("check whether table has existed error: {}", e.getMessage()); throw new RuntimeException(e); @@ -377,10 +395,9 @@ public static boolean checkTableExists(Connection connection, String tableName) stmt.close(); } } catch (SQLException e) { - log.error("close result or preparedStatement error: {}", e.getMessage()); + log.error("close result or statement error: {}", e.getMessage()); } } } - } diff --git a/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/converter/HiveParamConverterTest.java b/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/converter/HiveParamConverterTest.java index 961eb21..e469580 100644 --- a/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/converter/HiveParamConverterTest.java +++ b/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/converter/HiveParamConverterTest.java @@ -16,158 +16,158 @@ package org.secretflow.dataproxy.plugin.hive.converter; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.InjectMocks; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.secretflow.dataproxy.plugin.database.config.DatabaseConnectConfig; -import org.secretflow.dataproxy.plugin.database.config.DatabaseTableQueryConfig; -import org.secretflow.dataproxy.plugin.database.config.DatabaseWriteConfig; -import org.secretflow.dataproxy.plugin.database.config.ScqlCommandJobConfig; -import org.secretflow.dataproxy.plugin.database.constant.DatabaseTypeEnum; -import org.secretflow.dataproxy.plugin.database.converter.DatabaseParamConverter; -import org.secretflow.v1alpha1.kusciaapi.Domaindata; -import org.secretflow.v1alpha1.kusciaapi.Domaindatasource; -import org.secretflow.v1alpha1.kusciaapi.Flightdm; -import org.secretflow.v1alpha1.kusciaapi.Flightinner; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class HiveParamConverterTest { - @InjectMocks - private DatabaseParamConverter hiveParamConverter; - - @Mock - private Flightinner.CommandDataMeshSqlQuery meshSqlQuery; - - @Mock - private Flightinner.CommandDataMeshQuery meshQuery; - - @Mock - private Flightinner.CommandDataMeshUpdate meshUpdate; - - private final Domaindatasource.DatabaseDataSourceInfo hiveDataSourceInfo = - Domaindatasource.DatabaseDataSourceInfo - .newBuilder() - .setEndpoint("endpoint") - .setDatabase("database") - .setUser("user") - .setPassword("password") - .build(); - private final Domaindatasource.DataSourceInfo dataSourceInfo = - Domaindatasource.DataSourceInfo.newBuilder().setDatabase(hiveDataSourceInfo).build(); - - private final Domaindatasource.DomainDataSource domainDataSource = - Domaindatasource.DomainDataSource.newBuilder() - .setDatasourceId("datasourceId") - .setName("datasourceName") - .setType("hive") - .setInfo(dataSourceInfo) - .build(); - - private final Domaindata.DomainData domainData = - Domaindata.DomainData.newBuilder() - .setDatasourceId("datasourceId") - .setName("domainDataName") - .setRelativeUri("table_name_or_file_path") - .setDomaindataId("domainDataId") - .setType("table") - .build(); - private final Flightdm.CommandDomainDataQuery commandDomainDataQueryWithCsv = - Flightdm.CommandDomainDataQuery.newBuilder() - .setContentType(Flightdm.ContentType.CSV) - .setPartitionSpec("partition_spec") - .build(); - - private final Flightdm.CommandDomainDataUpdate commandDomainDataUpdateWithCsv = - Flightdm.CommandDomainDataUpdate.newBuilder() - .setContentType(Flightdm.ContentType.CSV) - .setPartitionSpec("partition_spec") - .build(); - /** - * Test Scenario: Test the convert method and enter the CommandDataMeshSqlQuery type - */ - @Test - public void testConvertMeshSqlQuery() { - String testSql = "select * from test_table;"; - final Flightdm.CommandDataSourceSqlQuery sqlQuery = Flightdm.CommandDataSourceSqlQuery.newBuilder() - .setDatasourceId("dataSourceId") - .setSql(testSql).build(); - when(meshSqlQuery.getQuery()).thenReturn(sqlQuery); - when(meshSqlQuery.getDatasource()).thenReturn(domainDataSource); - - ScqlCommandJobConfig result = hiveParamConverter.convert(meshSqlQuery); - assertNotNull(result); - assertEquals(DatabaseTypeEnum.SQL, result.getDbTypeEnum()); - assertEquals(testSql, result.taskRunSQL()); - - this.testHiveConnectConfig(result.getDbConnectConfig()); - } - - - /** - * Test scenario: Test the convert method, enter the CommandDataMeshQuery type, and the ContentType to CSV - */ - @Test - public void testConvertMeshQueryWithCsvContentType() { - when(meshQuery.getQuery()).thenReturn(commandDomainDataQueryWithCsv); - when(meshQuery.getDomaindata()).thenReturn(domainData); - when(meshQuery.getDatasource()).thenReturn(domainDataSource); - - DatabaseTableQueryConfig result = this.testConvertMeshQueryWithNonType(); - assertEquals(DatabaseTypeEnum.TABLE, result.getDbTypeEnum()); - } - - - /** - * Test Scenario: Test the convert method, enter the CommandDataMeshUpdate type, and the ContentType to CSV - */ - @Test - public void testConvertMeshUpdateWithCsvContentType() { - - when(meshUpdate.getUpdate()).thenReturn(commandDomainDataUpdateWithCsv); - when(meshUpdate.getDomaindata()).thenReturn(domainData); - when(meshUpdate.getDatasource()).thenReturn(domainDataSource); - - DatabaseWriteConfig odpsWriteConfig = this.testConvertMeshUpdateWithNonType(); - assertEquals(DatabaseTypeEnum.TABLE, odpsWriteConfig.getDbTypeEnum()); - } - - private void testHiveConnectConfig(DatabaseConnectConfig hiveConnectConfig) { - assertNotNull(hiveConnectConfig); - assertEquals("endpoint", hiveConnectConfig.endpoint()); - assertEquals("database", hiveConnectConfig.database()); - assertEquals("user", hiveConnectConfig.username()); - assertEquals("password", hiveConnectConfig.password()); - } - - private DatabaseTableQueryConfig testConvertMeshQueryWithNonType() { - DatabaseTableQueryConfig result = hiveParamConverter.convert(meshQuery); - assertNotNull(result); - - this.testHiveConnectConfig(result.getDbConnectConfig()); - - assertNotNull(result.getCommandConfig()); - assertEquals("table_name_or_file_path", result.getCommandConfig().tableName()); - assertEquals("partition_spec", result.getCommandConfig().partition()); - - return result; - } - - private DatabaseWriteConfig testConvertMeshUpdateWithNonType() { - DatabaseWriteConfig result = hiveParamConverter.convert(meshUpdate); - assertNotNull(result); - - this.testHiveConnectConfig(result.getDbConnectConfig()); - - assertNotNull(result.getCommandConfig()); - assertEquals("table_name_or_file_path", result.getCommandConfig().tableName()); - assertEquals("partition_spec", result.getCommandConfig().partition()); - - return result; - } -} +//import org.junit.jupiter.api.Test; +//import org.junit.jupiter.api.extension.ExtendWith; +//import org.mockito.InjectMocks; +//import org.mockito.Mock; +//import org.mockito.junit.jupiter.MockitoExtension; +//import org.secretflow.dataproxy.plugin.database.config.DatabaseConnectConfig; +//import org.secretflow.dataproxy.plugin.database.config.DatabaseTableQueryConfig; +//import org.secretflow.dataproxy.plugin.database.config.DatabaseWriteConfig; +//import org.secretflow.dataproxy.plugin.database.config.ScqlCommandJobConfig; +//import org.secretflow.dataproxy.plugin.database.constant.DatabaseTypeEnum; +//import org.secretflow.dataproxy.plugin.database.converter.DatabaseParamConverter; +//import org.secretflow.v1alpha1.kusciaapi.Domaindata; +//import org.secretflow.v1alpha1.kusciaapi.Domaindatasource; +//import org.secretflow.v1alpha1.kusciaapi.Flightdm; +//import org.secretflow.v1alpha1.kusciaapi.Flightinner; +// +//import static org.junit.jupiter.api.Assertions.assertEquals; +//import static org.junit.jupiter.api.Assertions.assertNotNull; +//import static org.mockito.Mockito.when; +// +//@ExtendWith(MockitoExtension.class) +//public class HiveParamConverterTest { +// @InjectMocks +// private DatabaseParamConverter hiveParamConverter; +// +// @Mock +// private Flightinner.CommandDataMeshSqlQuery meshSqlQuery; +// +// @Mock +// private Flightinner.CommandDataMeshQuery meshQuery; +// +// @Mock +// private Flightinner.CommandDataMeshUpdate meshUpdate; +// +// private final Domaindatasource.DatabaseDataSourceInfo hiveDataSourceInfo = +// Domaindatasource.DatabaseDataSourceInfo +// .newBuilder() +// .setEndpoint("endpoint") +// .setDatabase("database") +// .setUser("user") +// .setPassword("password") +// .build(); +// private final Domaindatasource.DataSourceInfo dataSourceInfo = +// Domaindatasource.DataSourceInfo.newBuilder().setDatabase(hiveDataSourceInfo).build(); +// +// private final Domaindatasource.DomainDataSource domainDataSource = +// Domaindatasource.DomainDataSource.newBuilder() +// .setDatasourceId("datasourceId") +// .setName("datasourceName") +// .setType("hive") +// .setInfo(dataSourceInfo) +// .build(); +// +// private final Domaindata.DomainData domainData = +// Domaindata.DomainData.newBuilder() +// .setDatasourceId("datasourceId") +// .setName("domainDataName") +// .setRelativeUri("table_name_or_file_path") +// .setDomaindataId("domainDataId") +// .setType("table") +// .build(); +// private final Flightdm.CommandDomainDataQuery commandDomainDataQueryWithCsv = +// Flightdm.CommandDomainDataQuery.newBuilder() +// .setContentType(Flightdm.ContentType.CSV) +// .setPartitionSpec("partition_spec") +// .build(); +// +// private final Flightdm.CommandDomainDataUpdate commandDomainDataUpdateWithCsv = +// Flightdm.CommandDomainDataUpdate.newBuilder() +// .setContentType(Flightdm.ContentType.CSV) +// .setPartitionSpec("partition_spec") +// .build(); +// /** +// * Test Scenario: Test the convert method and enter the CommandDataMeshSqlQuery type +// */ +// @Test +// public void testConvertMeshSqlQuery() { +// String testSql = "select * from test_table;"; +// final Flightdm.CommandDataSourceSqlQuery sqlQuery = Flightdm.CommandDataSourceSqlQuery.newBuilder() +// .setDatasourceId("dataSourceId") +// .setSql(testSql).build(); +// when(meshSqlQuery.getQuery()).thenReturn(sqlQuery); +// when(meshSqlQuery.getDatasource()).thenReturn(domainDataSource); +// +// ScqlCommandJobConfig result = hiveParamConverter.convert(meshSqlQuery); +// assertNotNull(result); +// assertEquals(DatabaseTypeEnum.SQL, result.getDbTypeEnum()); +// assertEquals(testSql, result.taskRunSQL()); +// +// this.testHiveConnectConfig(result.getDbConnectConfig()); +// } +// +// +// /** +// * Test scenario: Test the convert method, enter the CommandDataMeshQuery type, and the ContentType to CSV +// */ +// @Test +// public void testConvertMeshQueryWithCsvContentType() { +// when(meshQuery.getQuery()).thenReturn(commandDomainDataQueryWithCsv); +// when(meshQuery.getDomaindata()).thenReturn(domainData); +// when(meshQuery.getDatasource()).thenReturn(domainDataSource); +// +// DatabaseTableQueryConfig result = this.testConvertMeshQueryWithNonType(); +// assertEquals(DatabaseTypeEnum.TABLE, result.getDbTypeEnum()); +// } +// +// +// /** +// * Test Scenario: Test the convert method, enter the CommandDataMeshUpdate type, and the ContentType to CSV +// */ +// @Test +// public void testConvertMeshUpdateWithCsvContentType() { +// +// when(meshUpdate.getUpdate()).thenReturn(commandDomainDataUpdateWithCsv); +// when(meshUpdate.getDomaindata()).thenReturn(domainData); +// when(meshUpdate.getDatasource()).thenReturn(domainDataSource); +// +// DatabaseWriteConfig odpsWriteConfig = this.testConvertMeshUpdateWithNonType(); +// assertEquals(DatabaseTypeEnum.TABLE, odpsWriteConfig.getDbTypeEnum()); +// } +// +// private void testHiveConnectConfig(DatabaseConnectConfig hiveConnectConfig) { +// assertNotNull(hiveConnectConfig); +// assertEquals("endpoint", hiveConnectConfig.endpoint()); +// assertEquals("database", hiveConnectConfig.database()); +// assertEquals("user", hiveConnectConfig.username()); +// assertEquals("password", hiveConnectConfig.password()); +// } +// +// private DatabaseTableQueryConfig testConvertMeshQueryWithNonType() { +// DatabaseTableQueryConfig result = hiveParamConverter.convert(meshQuery); +// assertNotNull(result); +// +// this.testHiveConnectConfig(result.getDbConnectConfig()); +// +// assertNotNull(result.getCommandConfig()); +// assertEquals("table_name_or_file_path", result.getCommandConfig().tableName()); +// assertEquals("partition_spec", result.getCommandConfig().partition()); +// +// return result; +// } +// +// private DatabaseWriteConfig testConvertMeshUpdateWithNonType() { +// DatabaseWriteConfig result = hiveParamConverter.convert(meshUpdate); +// assertNotNull(result); +// +// this.testHiveConnectConfig(result.getDbConnectConfig()); +// +// assertNotNull(result.getCommandConfig()); +// assertEquals("table_name_or_file_path", result.getCommandConfig().tableName()); +// assertEquals("partition_spec", result.getCommandConfig().partition()); +// +// return result; +// } +//} diff --git a/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducerTest.java b/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducerTest.java index 82f5fb3..983ef3b 100644 --- a/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducerTest.java +++ b/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/producer/HiveFlightProducerTest.java @@ -16,106 +16,106 @@ package org.secretflow.dataproxy.plugin.hive.producer; -import com.google.protobuf.Any; -import org.apache.arrow.flight.FlightDescriptor; -import org.apache.arrow.flight.FlightInfo; -import org.apache.arrow.flight.FlightProducer; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.InjectMocks; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.secretflow.v1alpha1.kusciaapi.Domaindata; -import org.secretflow.v1alpha1.kusciaapi.Domaindatasource; -import org.secretflow.v1alpha1.kusciaapi.Flightdm; -import org.secretflow.v1alpha1.kusciaapi.Flightinner; - -import java.nio.charset.StandardCharsets; - -import static org.junit.jupiter.api.Assertions.*; -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.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class HiveFlightProducerTest { - @InjectMocks - private HiveFlightProducer hiveFlightProducer; - - @Mock - private FlightProducer.CallContext context; - - @Mock - private FlightDescriptor descriptor; - - private final Domaindatasource.DatabaseDataSourceInfo hiveDataSourceInfo = - Domaindatasource.DatabaseDataSourceInfo - .newBuilder() - .setEndpoint("endpoint") - .setDatabase("database") - .setUser("user") - .setPassword("password") - .build(); - private final Domaindatasource.DataSourceInfo dataSourceInfo = - Domaindatasource.DataSourceInfo.newBuilder().setDatabase(hiveDataSourceInfo).build(); - - private final Domaindatasource.DomainDataSource domainDataSource = - Domaindatasource.DomainDataSource.newBuilder() - .setDatasourceId("datasourceId") - .setName("datasourceName") - .setType("hive") - .setInfo(dataSourceInfo) - .build(); - - private final Domaindata.DomainData domainData = - Domaindata.DomainData.newBuilder() - .setDatasourceId("datasourceId") - .setName("domainDataName") - .setRelativeUri("table_name_or_file_path") - .setDomaindataId("domainDataId") - .setType("table") - .build(); - - private final Flightdm.CommandDomainDataQuery commandDomainDataQueryWithCSV = - Flightdm.CommandDomainDataQuery.newBuilder() - .setContentType(Flightdm.ContentType.CSV) - .setPartitionSpec("partition_spec") - .build(); - - @Test - public void testGetProducerName() { - String producerName = hiveFlightProducer.getProducerName(); - assertEquals("hive", producerName); - } - - @Test - public void testGetFlightInfoWithTableCommand() { - Flightinner.CommandDataMeshQuery dataMeshQuery = - Flightinner.CommandDataMeshQuery.newBuilder() - .setQuery(commandDomainDataQueryWithCSV) - .setDatasource(domainDataSource) - .setDomaindata(domainData) - .build(); - when(descriptor.getCommand()).thenReturn(Any.pack(dataMeshQuery).toByteArray()); - assertDoesNotThrow(() -> { - FlightInfo flightInfo = hiveFlightProducer.getFlightInfo(context, descriptor); - - assertNotNull(flightInfo); - assertFalse(flightInfo.getEndpoints().isEmpty()); - - assertNotNull(flightInfo.getEndpoints().get(0).getLocations()); - assertFalse(flightInfo.getEndpoints().get(0).getLocations().isEmpty()); - - assertNotNull(flightInfo.getEndpoints().get(0).getTicket()); - assertNotNull(flightInfo.getEndpoints().get(0).getTicket().getBytes()); - - }); - } - - @Test - public void testGetFlightInfoWithUnsupportedType() { - when(descriptor.getCommand()).thenReturn("testCommand".getBytes(StandardCharsets.UTF_8)); - assertThrows(RuntimeException.class, () -> hiveFlightProducer.getFlightInfo(context, descriptor)); - } -} +//import com.google.protobuf.Any; +//import org.apache.arrow.flight.FlightDescriptor; +//import org.apache.arrow.flight.FlightInfo; +//import org.apache.arrow.flight.FlightProducer; +//import org.junit.jupiter.api.Test; +//import org.junit.jupiter.api.extension.ExtendWith; +//import org.mockito.InjectMocks; +//import org.mockito.Mock; +//import org.mockito.junit.jupiter.MockitoExtension; +//import org.secretflow.v1alpha1.kusciaapi.Domaindata; +//import org.secretflow.v1alpha1.kusciaapi.Domaindatasource; +//import org.secretflow.v1alpha1.kusciaapi.Flightdm; +//import org.secretflow.v1alpha1.kusciaapi.Flightinner; +// +//import java.nio.charset.StandardCharsets; +// +//import static org.junit.jupiter.api.Assertions.*; +//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.mockito.Mockito.when; +// +//@ExtendWith(MockitoExtension.class) +//public class HiveFlightProducerTest { +// @InjectMocks +// private HiveFlightProducer hiveFlightProducer; +// +// @Mock +// private FlightProducer.CallContext context; +// +// @Mock +// private FlightDescriptor descriptor; +// +// private final Domaindatasource.DatabaseDataSourceInfo hiveDataSourceInfo = +// Domaindatasource.DatabaseDataSourceInfo +// .newBuilder() +// .setEndpoint("endpoint") +// .setDatabase("database") +// .setUser("user") +// .setPassword("password") +// .build(); +// private final Domaindatasource.DataSourceInfo dataSourceInfo = +// Domaindatasource.DataSourceInfo.newBuilder().setDatabase(hiveDataSourceInfo).build(); +// +// private final Domaindatasource.DomainDataSource domainDataSource = +// Domaindatasource.DomainDataSource.newBuilder() +// .setDatasourceId("datasourceId") +// .setName("datasourceName") +// .setType("hive") +// .setInfo(dataSourceInfo) +// .build(); +// +// private final Domaindata.DomainData domainData = +// Domaindata.DomainData.newBuilder() +// .setDatasourceId("datasourceId") +// .setName("domainDataName") +// .setRelativeUri("table_name_or_file_path") +// .setDomaindataId("domainDataId") +// .setType("table") +// .build(); +// +// private final Flightdm.CommandDomainDataQuery commandDomainDataQueryWithCSV = +// Flightdm.CommandDomainDataQuery.newBuilder() +// .setContentType(Flightdm.ContentType.CSV) +// .setPartitionSpec("partition_spec") +// .build(); +// +// @Test +// public void testGetProducerName() { +// String producerName = hiveFlightProducer.getProducerName(); +// assertEquals("hive", producerName); +// } +// +// @Test +// public void testGetFlightInfoWithTableCommand() { +// Flightinner.CommandDataMeshQuery dataMeshQuery = +// Flightinner.CommandDataMeshQuery.newBuilder() +// .setQuery(commandDomainDataQueryWithCSV) +// .setDatasource(domainDataSource) +// .setDomaindata(domainData) +// .build(); +// when(descriptor.getCommand()).thenReturn(Any.pack(dataMeshQuery).toByteArray()); +// assertDoesNotThrow(() -> { +// FlightInfo flightInfo = hiveFlightProducer.getFlightInfo(context, descriptor); +// +// assertNotNull(flightInfo); +// assertFalse(flightInfo.getEndpoints().isEmpty()); +// +// assertNotNull(flightInfo.getEndpoints().get(0).getLocations()); +// assertFalse(flightInfo.getEndpoints().get(0).getLocations().isEmpty()); +// +// assertNotNull(flightInfo.getEndpoints().get(0).getTicket()); +// assertNotNull(flightInfo.getEndpoints().get(0).getTicket().getBytes()); +// +// }); +// } +// +// @Test +// public void testGetFlightInfoWithUnsupportedType() { +// when(descriptor.getCommand()).thenReturn("testCommand".getBytes(StandardCharsets.UTF_8)); +// assertThrows(RuntimeException.class, () -> hiveFlightProducer.getFlightInfo(context, descriptor)); +// } +//} diff --git a/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtilTest.java b/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtilTest.java index 1e31b0b..88dd726 100644 --- a/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtilTest.java +++ b/dataproxy-plugins/dataproxy-plugin-hive/src/test/java/org/secretflow/dataproxy/plugin/hive/utils/HiveUtilTest.java @@ -17,30 +17,44 @@ package org.secretflow.dataproxy.plugin.hive.utils; import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; + +import org.secretflow.dataproxy.common.exceptions.DataproxyException; + import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; +import org.secretflow.dataproxy.plugin.database.writer.DatabaseRecordWriter; + import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.Map; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.secretflow.dataproxy.plugin.hive.utils.HiveUtil.buildMultiRowInsertSql; @ExtendWith(MockitoExtension.class) -public class HiveUtilTest { +class HiveUtilTest { - @Test - public void testBuildQuerySql() { - String tableName = "table"; - List fields = new ArrayList<>(); - fields.add("col1"); - fields.add("col2"); - fields.add("col3"); - String stmt = HiveUtil.buildQuerySql(tableName, fields, "region=us, date=2025-06-26"); - assertEquals("select col1,col2,col3 from table where region='us' and date='2025-06-26'", stmt); - } +// @Test +// void testBuildQuerySql() { +// String tableName = "table"; +// List fields = new ArrayList<>(); +// fields.add("col1"); +// fields.add("col2"); +// fields.add("col3"); +// String stmt = HiveUtil.buildQuerySql(tableName, fields, "region=us, date=2025-06-26"); +// assertEquals("select col1,col2,col3 from table where region='us' and date='2025-06-26'", stmt); +// } @Test - public void testJDBCType2ArrowType() { + void testJDBCType2ArrowType() { assertEquals(Types.MinorType.INT.getType(), HiveUtil.jdbcType2ArrowType("int")); assertEquals(Types.MinorType.BIGINT.getType(), HiveUtil.jdbcType2ArrowType("bigint")); assertEquals(Types.MinorType.FLOAT4.getType(), HiveUtil.jdbcType2ArrowType("float")); @@ -51,4 +65,87 @@ public void testJDBCType2ArrowType() { assertEquals(Types.MinorType.TIMESTAMPMILLI.getType(), HiveUtil.jdbcType2ArrowType("timestamp")); } + @Test + void testNormal() { + List fields = Arrays.asList( + Field.nullable("id", new ArrowType.Int(32, true)), + Field.nullable("name", new ArrowType.Utf8()), + Field.nullable("dt", new ArrowType.Utf8()) + ); + Schema schema = new Schema(fields); + + List> data = List.of( + Map.of("id", 1, "name", "Alice"), + Map.of("id", 2, "name", "Bob") + ); + Map partition = Map.of("dt", "2024-06-01"); + + DatabaseRecordWriter.SqlWithParams sp = buildMultiRowInsertSql("user", schema, data, partition); + + // SQL 模板校验 + assertEquals( + "INSERT INTO TABLE user PARTITION (dt = ?) (id, name) VALUES (?, ?), (?, ?)", + sp.sql + ); + + // Parameter validation: partition values + row values + assertEquals(List.of("2024-06-01", 1, "Alice", 2, "Bob"), sp.params); + } + + /** + * Throws exception when no data is provided + */ + @Test + void testEmptyData() { + Schema schema = new Schema(Collections.emptyList()); + Exception ex = assertThrows( + IllegalArgumentException.class, + () -> buildMultiRowInsertSql("t", schema, Collections.emptyList(), Collections.emptyMap()) + ); + assertEquals("No data to insert", ex.getMessage()); + } + + /** + * Throws exception when table name is invalid + */ + @Test + void testInvalidTableName() { + Schema schema = new Schema(List.of(Field.nullable("a", new ArrowType.Int(32, true)))); + Exception ex = assertThrows( + DataproxyException.class, + () -> buildMultiRowInsertSql("user;--", schema, + List.of(Map.of("a", 1)), Collections.emptyMap()) + ); + assertTrue(ex.getMessage().contains("Invalid tableName")); + } + + /** + * Throws exception when field name is invalid + */ + @Test + void testInvalidFieldName() { + List fields = List.of(Field.nullable("a b", new ArrowType.Int(32, true))); + Schema schema = new Schema(fields); + Exception ex = assertThrows( + DataproxyException.class, + () -> buildMultiRowInsertSql("t", schema, + List.of(Map.of("a b", 1)), Collections.emptyMap()) + ); + assertTrue(ex.getMessage().contains("Invalid field name")); + } + + /** + * Throws exception when partition key is invalid + */ + @Test + void testInvalidPartitionKey() { + Schema schema = new Schema(List.of(Field.nullable("a", new ArrowType.Int(32, true)))); + Exception ex = assertThrows( + DataproxyException.class, + () -> buildMultiRowInsertSql("t", schema, + List.of(Map.of("a", 1)), Map.of("a b", "v")) + ); + assertTrue(ex.getMessage().contains("Invalid partition key")); + } + } diff --git a/dataproxy-plugins/dataproxy-plugin-odps/src/main/java/org/secretflow/dataproxy/plugin/odps/producer/OdpsFlightProducer.java b/dataproxy-plugins/dataproxy-plugin-odps/src/main/java/org/secretflow/dataproxy/plugin/odps/producer/OdpsFlightProducer.java index d694142..7c6edc2 100644 --- a/dataproxy-plugins/dataproxy-plugin-odps/src/main/java/org/secretflow/dataproxy/plugin/odps/producer/OdpsFlightProducer.java +++ b/dataproxy-plugins/dataproxy-plugin-odps/src/main/java/org/secretflow/dataproxy/plugin/odps/producer/OdpsFlightProducer.java @@ -69,6 +69,7 @@ public class OdpsFlightProducer extends NoOpFlightProducer implements DataProxyFlightProducer { private final TicketService ticketService = CacheTicketService.getInstance(); + /** * Obtain the data type used for registration name and identification processing. * @@ -119,7 +120,7 @@ public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor } // Only the protocol is used, and the concrete schema is not returned here. - return new FlightInfo(DataProxyFlightProducer.DEFACT_SCHEMA, descriptor, endpointList, 0, 0,true, IpcOption.DEFAULT); + return new FlightInfo(DataProxyFlightProducer.DEFACT_SCHEMA, descriptor, endpointList, 0, 0, true, IpcOption.DEFAULT); } catch (InvalidProtocolBufferException e) { throw CallStatus.INVALID_ARGUMENT .withCause(e) @@ -136,8 +137,11 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l ParamWrapper paramWrapper = ticketService.getParamWrapper(ticket.getBytes()); ArrowReader odpsReader = null; + BufferAllocator allocator = null; try { - + allocator = + FlightServerContext.getInstance().getFlightServerConfig().getBufferAllocator() + .newChildAllocator("odpsReader", 0, 2 * 128 * 1024 * 1024); Object param = paramWrapper.param(); if (param == null) { @@ -148,7 +152,7 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l if (OdpsTypeEnum.FILE.equals(odpsCommandConfig.getOdpsTypeEnum())) { Object commandConfig = odpsCommandConfig.getCommandConfig(); if (commandConfig instanceof OdpsTableConfig odpsTableConfig) { - odpsReader = new OdpsResourceReader(new RootAllocator(), odpsCommandConfig.getOdpsConnectConfig(), odpsTableConfig); + odpsReader = new OdpsResourceReader(allocator, odpsCommandConfig.getOdpsConnectConfig(), odpsTableConfig); } else { throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "The odps read parameter is invalid, type url: " + commandConfig.getClass()); } @@ -163,13 +167,14 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l odpsReader = new OdpsReader(new RootAllocator(), taskConfigs.get(0)); } } else if (param instanceof TaskConfig taskConfig) { - odpsReader = new OdpsReader(new RootAllocator(), taskConfig); + odpsReader = new OdpsReader(allocator, taskConfig); } else { throw DataproxyException.of(DataproxyErrorCode.PARAMS_UNRELIABLE, "The odps read parameter is invalid, type url: " + param.getClass()); } listener.start(odpsReader.getVectorSchemaRoot()); while (true) { + log.debug("Before put AllocatedMemory: {}", allocator.getAllocatedMemory()); if (context.isCancelled()) { log.warn("reader is cancelled"); break; @@ -177,6 +182,7 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l if (odpsReader.loadNextBatch()) { listener.putNext(); + log.debug("After put AllocatedMemory: {}", allocator.getAllocatedMemory()); } else { break; } @@ -191,9 +197,14 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l .toRuntimeException(); } finally { try { + log.info("close odps reader"); + if (odpsReader != null) { odpsReader.close(); } + if (allocator != null) { + allocator.close(); + } } catch (Exception e) { log.error("close odps reader error", e); } @@ -231,8 +242,8 @@ public Runnable acceptPut( askMsg = "row count: " + rowCount; writer.write(vectorSchemaRoot); - try(BufferAllocator ba = new RootAllocator(1024); - final ArrowBuf buffer = ba.buffer(askMsg.getBytes(StandardCharsets.UTF_8).length)) { + try (BufferAllocator ba = new RootAllocator(1024); + final ArrowBuf buffer = ba.buffer(askMsg.getBytes(StandardCharsets.UTF_8).length)) { ackStream.onNext(PutResult.metadata(buffer)); } count += rowCount; diff --git a/dataproxy-server/pom.xml b/dataproxy-server/pom.xml index 5aad838..c233c51 100644 --- a/dataproxy-server/pom.xml +++ b/dataproxy-server/pom.xml @@ -40,6 +40,11 @@ dataproxy-plugin-odps + + org.secretflow + dataproxy-metrics + + org.secretflow dataproxy-plugin-database @@ -70,6 +75,30 @@ lombok + + + org.junit.jupiter + junit-jupiter-api + test + + + org.junit.jupiter + junit-jupiter-engine + test + + + org.mockito + mockito-inline + test + + + + + org.mockito + mockito-junit-jupiter + test + + @@ -94,4 +123,4 @@ - \ No newline at end of file + diff --git a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyFlightServer.java b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyFlightServer.java index 883f502..3b536d8 100644 --- a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyFlightServer.java +++ b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyFlightServer.java @@ -18,8 +18,6 @@ import lombok.extern.slf4j.Slf4j; import org.apache.arrow.flight.FlightServer; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; import org.secretflow.dataproxy.core.config.FlightServerConfig; import org.secretflow.dataproxy.core.spi.producer.DataProxyFlightProducer; import org.secretflow.dataproxy.server.flight.CompositeFlightProducer; @@ -59,14 +57,12 @@ public void close() throws Exception { } private FlightServer init(FlightServerConfig config) { - - BufferAllocator allocator = new RootAllocator(); - + log.info("DataProxyFlightServer init."); return FlightServer.builder() // .useTls(null, null) // .useMTlsClientVerification(null) .middleware(FlightServerTraceMiddleware.getKey(), new FlightServerTraceMiddleware.FlightServerTraceMiddlewareFactory()) - .allocator(allocator) + .allocator(config.getBufferAllocator()) .location(config.getLocation()) .producer(initProducer()) .build(); diff --git a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyMetricsServer.java b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyMetricsServer.java new file mode 100644 index 0000000..dc0dddb --- /dev/null +++ b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyMetricsServer.java @@ -0,0 +1,100 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server; + +import io.netty.bootstrap.ServerBootstrap; +import io.netty.channel.Channel; +import io.netty.channel.EventLoopGroup; +import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.socket.nio.NioServerSocketChannel; +import io.netty.handler.logging.LogLevel; +import io.netty.handler.logging.LoggingHandler; +import lombok.extern.slf4j.Slf4j; +import org.secretflow.dataproxy.core.config.FlightServerConfigKey; +import org.secretflow.dataproxy.core.config.FlightServerContext; +import org.secretflow.dataproxy.server.handler.MetricsChannelHandler; + +import java.util.function.Supplier; + +/** + * @author yuexie + * @date 2025/4/11 18:06 + **/ +@Slf4j +public class DataProxyMetricsServer { + + public static final int DEFAULT_METRICS_PORT = 9101; + private static final String METRICS_ENDPOINT = "/metrics"; + + private final EventLoopGroup bossGroup; + private final EventLoopGroup workerGroup; + private final Supplier hostAddressSupplier; + + public DataProxyMetricsServer() { + this.bossGroup = new NioEventLoopGroup(1); + this.workerGroup = new NioEventLoopGroup(); + this.hostAddressSupplier = this::getHostAddress; + } + + // Constructor for testing + public DataProxyMetricsServer(EventLoopGroup bossGroup, EventLoopGroup workerGroup, Supplier hostAddressSupplier) { + this.bossGroup = bossGroup; + this.workerGroup = workerGroup; + this.hostAddressSupplier = hostAddressSupplier; + } + + public void start() throws InterruptedException { + Integer metricsPort = FlightServerContext.getOrDefault(FlightServerConfigKey.METRICS_PORT, Integer.class, DEFAULT_METRICS_PORT); + log.info("Starting Metrics Server on port {}", metricsPort); + + try { + ServerBootstrap b = createServerBootstrap(); + Channel ch = b.bind(DEFAULT_METRICS_PORT).sync().channel(); + log.info("Metrics server started at http://{}:{}{}", + hostAddressSupplier.get(), DEFAULT_METRICS_PORT, METRICS_ENDPOINT); + ch.closeFuture().sync(); + } catch (InterruptedException e) { + log.error("Metrics server was interrupted", e); + throw e; + } finally { + shutdownEventLoopGroups(bossGroup, workerGroup); + } + } + + private ServerBootstrap createServerBootstrap() { + ServerBootstrap b = new ServerBootstrap(); + b.group(bossGroup, workerGroup) + .channel(NioServerSocketChannel.class) + .handler(new LoggingHandler(LogLevel.INFO)) + .childHandler(new MetricsChannelHandler()); + return b; + } + + private void shutdownEventLoopGroups(EventLoopGroup bossGroup, EventLoopGroup workerGroup) { + try { + bossGroup.shutdownGracefully().sync(); + workerGroup.shutdownGracefully().sync(); + } catch (InterruptedException e) { + log.warn("Interrupted while shutting down event loop groups", e); + Thread.currentThread().interrupt(); + } + } + + String getHostAddress() { + return FlightServerContext.getInstance().getFlightServerConfig().host(); + } +} diff --git a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyServerApplication.java b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyServerApplication.java index 5241f44..0d23f90 100644 --- a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyServerApplication.java +++ b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/DataProxyServerApplication.java @@ -35,12 +35,22 @@ public static void main(String[] args) { try (DataProxyFlightServer dataProxyFlightServer = new DataProxyFlightServer(flightServerConfig)) { dataProxyFlightServer.start(); log.info("Data proxy flight server start at {}:{}", flightServerConfig.getLocation().getUri().getHost(), flightServerConfig.port()); + + new Thread(() -> { + try { + new DataProxyMetricsServer().start(); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }, "mertics-thread").start(); + dataProxyFlightServer.awaitTermination(); } catch (Exception e) { log.error("DataProxyFlightServer start failed", e); throw new RuntimeException(e); } finally { log.warn("DataProxyFlightServer stopped"); + flightServerConfig.getBufferAllocator().close(); } } } diff --git a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducer.java b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducer.java index c717653..c4839c1 100644 --- a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducer.java +++ b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducer.java @@ -241,9 +241,12 @@ private FlightProducer getProducer(FlightDescriptor descriptor) { yield CacheTicketService.getInstance().getParamWrapper(domaindataHandle.getBytes()).producerKey(); } - default -> throw CallStatus.INVALID_ARGUMENT + default -> { + log.error("Unknown command type: {}", any.getTypeUrl()); + throw CallStatus.INVALID_ARGUMENT .withDescription("Unknown command type") .toRuntimeException(); + } }; log.info("producer type is {}", dataSourceType); return registry.getOrDefaultNoOp(dataSourceType); diff --git a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/handler/MetricsChannelHandler.java b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/handler/MetricsChannelHandler.java new file mode 100644 index 0000000..aa7a3e5 --- /dev/null +++ b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/handler/MetricsChannelHandler.java @@ -0,0 +1,36 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server.handler; + +import io.netty.channel.ChannelInitializer; +import io.netty.channel.socket.SocketChannel; +import io.netty.handler.codec.http.HttpObjectAggregator; +import io.netty.handler.codec.http.HttpServerCodec; + +/** + * @author yuexie + * @date 2025/6/3 14:34 + **/ +public class MetricsChannelHandler extends ChannelInitializer { + + @Override + public void initChannel(SocketChannel ch) { + ch.pipeline().addLast(new HttpServerCodec()); + ch.pipeline().addLast(new HttpObjectAggregator(65536)); + ch.pipeline().addLast(new MetricsChannelInboundHandler()); + } +} diff --git a/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/handler/MetricsChannelInboundHandler.java b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/handler/MetricsChannelInboundHandler.java new file mode 100644 index 0000000..cde2f51 --- /dev/null +++ b/dataproxy-server/src/main/java/org/secretflow/dataproxy/server/handler/MetricsChannelInboundHandler.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server.handler; + +import io.micrometer.prometheusmetrics.PrometheusConfig; +import io.micrometer.prometheusmetrics.PrometheusMeterRegistry; +import io.netty.buffer.Unpooled; +import io.netty.channel.ChannelFutureListener; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.SimpleChannelInboundHandler; +import io.netty.handler.codec.http.DefaultFullHttpResponse; +import io.netty.handler.codec.http.FullHttpRequest; +import io.netty.handler.codec.http.FullHttpResponse; +import io.netty.handler.codec.http.HttpHeaderNames; +import io.netty.handler.codec.http.HttpResponseStatus; +import io.netty.handler.codec.http.HttpVersion; +import org.secretflow.dataproxy.metrics.JvmMetricsRegistrar; + +import java.nio.charset.StandardCharsets; + +/** + * @author yuexie + * @date 2025/6/3 14:40 + **/ +public class MetricsChannelInboundHandler extends SimpleChannelInboundHandler { + + private final static PrometheusMeterRegistry PROMETHEUS_METER_REGISTRY = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); + + static { + JvmMetricsRegistrar.registerClassLoaderMetrics(PROMETHEUS_METER_REGISTRY); + JvmMetricsRegistrar.registerJvmGcMetrics(PROMETHEUS_METER_REGISTRY); + JvmMetricsRegistrar.registerJvmMemoryMetrics(PROMETHEUS_METER_REGISTRY); + JvmMetricsRegistrar.registerJvmThreadMetrics(PROMETHEUS_METER_REGISTRY); + JvmMetricsRegistrar.registerFileDescriptorMetrics(PROMETHEUS_METER_REGISTRY); + } + + @Override + public void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) throws Exception { + if ("/metrics".equals(req.uri())) { + + String metrics = PROMETHEUS_METER_REGISTRY.scrape(); + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.OK, + Unpooled.copiedBuffer(metrics, StandardCharsets.UTF_8) + ); + response.headers().set(HttpHeaderNames.CONTENT_TYPE, "text/plain; version=0.0.4;charset=utf-8"); + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + + } else { + + FullHttpResponse response = new DefaultFullHttpResponse( + HttpVersion.HTTP_1_1, + HttpResponseStatus.NOT_FOUND + ); + ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); + + } + } +} diff --git a/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/DataProxyFlightServerTest.java b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/DataProxyFlightServerTest.java new file mode 100644 index 0000000..0d22b09 --- /dev/null +++ b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/DataProxyFlightServerTest.java @@ -0,0 +1,47 @@ +/* + * Copyright 2024 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server; + +import org.junit.jupiter.api.Test; +import org.secretflow.dataproxy.core.config.FlightServerConfig; +import org.secretflow.dataproxy.core.config.FlightServerContext; + +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; + +public class DataProxyFlightServerTest { + + @Test + void testStart() { + assertDoesNotThrow(() -> { + Thread thread = new Thread(() -> { + FlightServerConfig flightServerConfig = FlightServerContext.getInstance().getFlightServerConfig(); + try (DataProxyFlightServer server = new DataProxyFlightServer(flightServerConfig)) { + server.start(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + thread.start(); + TimeUnit.SECONDS.sleep(1); + thread.interrupt(); + }); + } +} diff --git a/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/DataProxyMetricsServerTest.java b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/DataProxyMetricsServerTest.java new file mode 100644 index 0000000..de3d030 --- /dev/null +++ b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/DataProxyMetricsServerTest.java @@ -0,0 +1,46 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; + +@ExtendWith(MockitoExtension.class) +public class DataProxyMetricsServerTest { + + @Test + void testStart() { + assertDoesNotThrow(() -> { + Thread thread = new Thread(() -> { + try { + new DataProxyMetricsServer().start(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + thread.start(); + TimeUnit.SECONDS.sleep(1); + thread.interrupt(); + }); + } + +} diff --git a/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducerTest.java b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducerTest.java new file mode 100644 index 0000000..55c2007 --- /dev/null +++ b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/flight/CompositeFlightProducerTest.java @@ -0,0 +1,215 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server.flight; + +import com.google.protobuf.Any; +import org.apache.arrow.flight.Criteria; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightProducer; +import org.apache.arrow.flight.FlightProducer.CallContext; +import org.apache.arrow.flight.FlightProducer.ServerStreamListener; +import org.apache.arrow.flight.SchemaResult; +import org.apache.arrow.flight.Ticket; +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 org.secretflow.dataproxy.core.param.ParamWrapper; +import org.secretflow.dataproxy.core.service.impl.CacheTicketService; +import org.secretflow.v1alpha1.kusciaapi.Domaindatasource; +import org.secretflow.v1alpha1.kusciaapi.Flightdm; +import org.secretflow.v1alpha1.kusciaapi.Flightinner; + +import java.nio.charset.StandardCharsets; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * @author yuexie + * @date 2025/6/9 14:45 + **/ +@ExtendWith(MockitoExtension.class) +public class CompositeFlightProducerTest { + + @Mock + private ProducerRegistry mockRegistry; + + @Mock + private FlightProducer mockFlightProducer; + + private CompositeFlightProducer producer; + + @BeforeEach + void setUp() { + producer = new CompositeFlightProducer(mockRegistry); + } + + @Test + void getStream() { + CallContext context = mock(CallContext.class); + Ticket ticket = new Ticket("test".getBytes()); + ServerStreamListener listener = mock(ServerStreamListener.class); + + when(mockRegistry.getOrDefaultNoOp(anyString())).thenReturn(mockFlightProducer); + + producer.getStream(context, ticket, listener); + + verify(mockFlightProducer).getStream(context, ticket, listener); + } + + @Test + void getStreamWithNullTicket() { + CallContext context = mock(CallContext.class); + ServerStreamListener listener = mock(ServerStreamListener.class); + + assertThrows(NullPointerException.class, () -> producer.getStream(context, null, listener)); + } + + @Test + void listFlights() { + CallContext context = mock(CallContext.class); + Criteria criteria = mock(Criteria.class); + FlightProducer.StreamListener listener = mock(FlightProducer.StreamListener.class); + + when(mockRegistry.getOrDefaultNoOp(anyString())).thenReturn(mockFlightProducer); + + producer.listFlights(context, criteria, listener); + + verify(mockFlightProducer).listFlights(context, criteria, listener); + } + + @Test + void listFlightsWithNullCriteria() { + CallContext context = mock(CallContext.class); + FlightProducer.StreamListener listener = mock(FlightProducer.StreamListener.class); + + assertThrows(NullPointerException.class, () -> producer.listFlights(context, null, listener)); + } + + + @Test + void getFlightInfoWithCommandDataMeshQuery() { + assertDoesNotThrow(() -> returnFlightInfoFromProducer(Any.pack(mockCommandDataMeshQuery()))); + } + + @Test + void getFlightInfoWithCommandDataMeshSqlQuery() { + assertDoesNotThrow(() -> returnFlightInfoFromProducer(Any.pack(mockCommandDataMeshSqlQuery()))); + } + + @Test + void getFlightInfoWithCommandDataMeshUpdate() { + assertDoesNotThrow(() -> returnFlightInfoFromProducer(Any.pack(mockCommandDataMeshUpdate()))); + } + + @Test + void getFlightInfoWithTicketDomainDataQuery() { + assertDoesNotThrow(() -> { + + ParamWrapper paramWrapper = ParamWrapper.of("mock", null); + + byte[] mocks = CacheTicketService.getInstance().generateTicket(paramWrapper); + + Flightdm.TicketDomainDataQuery ticketDomainDataQuery = mockTicketDomainDataQuery(new String(mocks, StandardCharsets.UTF_8)); + + paramWrapper.setParamIfAbsent(ticketDomainDataQuery); + returnFlightInfoFromProducer(Any.pack(ticketDomainDataQuery)); + }); + } + + + private void returnFlightInfoFromProducer(Any any) { + FlightDescriptor descriptor = FlightDescriptor.command(any.toByteArray()); + CallContext context = mock(CallContext.class); + FlightInfo expectedInfo = mock(FlightInfo.class); + + when(mockRegistry.getOrDefaultNoOp(anyString())).thenReturn(mockFlightProducer); + when(mockFlightProducer.getFlightInfo(context, descriptor)).thenReturn(expectedInfo); + + FlightInfo result = producer.getFlightInfo(context, descriptor); + + assertEquals(expectedInfo, result); + verify(mockFlightProducer).getFlightInfo(context, descriptor); + } + + @Test + void getSchemaWhenCommandDataMeshQuery() { + CallContext context = mock(CallContext.class); + + FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(mockCommandDataMeshQuery()).toByteArray()); + + SchemaResult mock = mock(SchemaResult.class); + + when(mockRegistry.getOrDefaultNoOp(anyString())).thenReturn(mockFlightProducer); + when(mockFlightProducer.getSchema(context, descriptor)).thenReturn(mock); + + SchemaResult schema = producer.getSchema(context, descriptor); + + verify(mockFlightProducer).getSchema(context, descriptor); + assertEquals(mock, schema); + + } + + @Test + void shouldHandleExceptionWhenGetSchema() { + + CallContext context = mock(CallContext.class); + FlightDescriptor descriptor = FlightDescriptor.command(Any.pack(mockCommandDataMeshQuery()).toByteArray()); + + SchemaResult mock = mock(SchemaResult.class); + + when(mockRegistry.getOrDefaultNoOp(anyString())).thenReturn(mockFlightProducer); + when(mockFlightProducer.getSchema(context, descriptor)).thenReturn(mock); + when(mockFlightProducer.getSchema(context, descriptor)).thenThrow(new RuntimeException("test error")); + + assertThrows(RuntimeException.class, () -> producer.getSchema(context, descriptor)); + verify(mockFlightProducer).getSchema(context, descriptor); + } + + private Flightinner.CommandDataMeshQuery mockCommandDataMeshQuery() { + return Flightinner.CommandDataMeshQuery.newBuilder() + .setDatasource(Domaindatasource.DomainDataSource.newBuilder().setType("mock")) + .build(); + } + + private Flightinner.CommandDataMeshUpdate mockCommandDataMeshUpdate() { + return Flightinner.CommandDataMeshUpdate.newBuilder() + .setDatasource(Domaindatasource.DomainDataSource.newBuilder().setType("mock")) + .build(); + } + + private Flightinner.CommandDataMeshSqlQuery mockCommandDataMeshSqlQuery() { + return Flightinner.CommandDataMeshSqlQuery.newBuilder() + .setDatasource(Domaindatasource.DomainDataSource.newBuilder().setType("mock")) + .build(); + } + + private Flightdm.TicketDomainDataQuery mockTicketDomainDataQuery(String domainDataHandle) { + + return Flightdm.TicketDomainDataQuery.newBuilder() + .setDomaindataHandle(domainDataHandle) + .build(); + } +} diff --git a/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/flight/ProducerRegistryTest.java b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/flight/ProducerRegistryTest.java new file mode 100644 index 0000000..9f3bea3 --- /dev/null +++ b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/flight/ProducerRegistryTest.java @@ -0,0 +1,124 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server.flight; + +import org.apache.arrow.flight.FlightProducer; +import org.apache.arrow.flight.NoOpFlightProducer; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.Mockito.mock; + +public class ProducerRegistryTest { + + private ProducerRegistry registry; + private FlightProducer mockProducer; + + @BeforeEach + void setUp() { + registry = ProducerRegistry.getInstance(); + mockProducer = mock(FlightProducer.class); + } + + @Test + void shouldReturnSameInstance() { + ProducerRegistry anotherInstance = ProducerRegistry.getInstance(); + assertSame(registry, anotherInstance); + } + + @Test + void shouldRegisterAndGetProducer() { + registry.register("testKey", mockProducer); + FlightProducer result = registry.getOrDefaultNoOp("testKey"); + + assertSame(mockProducer, result); + } + + @Test + void shouldReturnDefaultForUnknownKey() { + FlightProducer result = registry.getOrDefaultNoOp("unknown"); + + assertNotNull(result); + assertInstanceOf(NoOpFlightProducer.class, result); + } + + @Test + void shouldOverrideExistingProducer() { + registry.register("testKey", mockProducer); + + FlightProducer anotherProducer = mock(FlightProducer.class); + registry.register("testKey", anotherProducer); + + assertSame(anotherProducer, registry.getOrDefaultNoOp("testKey")); + } + + @Test + void shouldBeThreadSafeForConcurrentAccess() throws InterruptedException { + final int threadCount = 10; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch latch = new CountDownLatch(threadCount); + + for (int i = 0; i < threadCount; i++) { + final int index = i; + executor.execute(() -> { + FlightProducer producer = mock(FlightProducer.class); + registry.register("key" + index, producer); + latch.countDown(); + }); + } + + latch.await(5, TimeUnit.SECONDS); + + for (int i = 0; i < threadCount; i++) { + String key = "key" + i; + FlightProducer producer = registry.getOrDefaultNoOp(key); + assertNotNull(producer, "Producer for " + key + " should not be null"); + assertFalse(producer instanceof NoOpFlightProducer, "Should not return NoOpProducer for " + key); + } + } + + @Test + void shouldNotLoseRegistrationUnderConcurrency() throws InterruptedException { + final int threadCount = 100; + final String sharedKey = "sharedKey"; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch latch = new CountDownLatch(threadCount); + + for (int i = 0; i < threadCount; i++) { + executor.execute(() -> { + FlightProducer producer = mock(FlightProducer.class); + registry.register(sharedKey, producer); + latch.countDown(); + }); + } + + latch.await(5, TimeUnit.SECONDS); + + FlightProducer result = registry.getOrDefaultNoOp(sharedKey); + assertNotNull(result); + assertFalse(result instanceof NoOpFlightProducer); + } +} diff --git a/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/handler/MetricsChannelHandlerTest.java b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/handler/MetricsChannelHandlerTest.java new file mode 100644 index 0000000..750a025 --- /dev/null +++ b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/handler/MetricsChannelHandlerTest.java @@ -0,0 +1,106 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server.handler; + +import io.netty.channel.ChannelHandler; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.socket.SocketChannel; +import io.netty.handler.codec.http.HttpObjectAggregator; +import io.netty.handler.codec.http.HttpServerCodec; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.InOrder; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * @author yuexie + * @date 2025/6/5 14:14 + **/ + +@ExtendWith(MockitoExtension.class) +public class MetricsChannelHandlerTest { + + @Mock + private SocketChannel socketChannel; + @Mock + private ChannelPipeline pipeline; + + private final MetricsChannelHandler handler = new MetricsChannelHandler(); + + // Test Scenario: Verify the order in which processors are added when the channel is initialized + @Test + public void testInitChannel_OrderOfHandlers() { + when(socketChannel.pipeline()).thenReturn(pipeline); + + handler.initChannel(socketChannel); + + InOrder inOrder = inOrder(pipeline); + inOrder.verify(pipeline).addLast(any(HttpServerCodec.class)); + inOrder.verify(pipeline).addLast(any(HttpObjectAggregator.class)); + inOrder.verify(pipeline).addLast(any(MetricsChannelInboundHandler.class)); + } + + // Test Case: Verify the correctness of the HttpObjectAggregator parameters + @Test + public void testHttpObjectAggregator_MaxContentLength() { + when(socketChannel.pipeline()).thenReturn(pipeline); + ArgumentCaptor captor = ArgumentCaptor.forClass(HttpObjectAggregator.class); + handler.initChannel(socketChannel); + + verify(pipeline, times(3)).addLast(captor.capture()); + + assertEquals(3, captor.getAllValues().size()); + HttpObjectAggregator httpObjectAggregator = captor.getAllValues().get(1); + assertEquals(HttpObjectAggregator.class, httpObjectAggregator.getClass()); + assertEquals(65536, httpObjectAggregator.maxContentLength()); + } + + // [Single Test Case] Test Scenario: Verify that three processors are added correctly + @Test + public void testInitChannel_AllHandlersAdded() { + when(socketChannel.pipeline()).thenReturn(pipeline); + + handler.initChannel(socketChannel); + + verify(pipeline, times(3)).addLast(any()); + } + + // [Single Test Case] Test Scenario: Verify the correctness of the processor type + @Test + public void testHandlerTypes() { + when(socketChannel.pipeline()).thenReturn(pipeline); + ArgumentCaptor captor = ArgumentCaptor.forClass(Object.class); + + handler.initChannel(socketChannel); + + verify(pipeline, times(3)).addLast((ChannelHandler) captor.capture()); + assertEquals(HttpServerCodec.class, captor.getAllValues().get(0).getClass()); + assertEquals(HttpObjectAggregator.class, captor.getAllValues().get(1).getClass()); + assertEquals(MetricsChannelInboundHandler.class, captor.getAllValues().get(2).getClass()); + } + + +} diff --git a/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/handler/MetricsChannelInboundHandlerTest.java b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/handler/MetricsChannelInboundHandlerTest.java new file mode 100644 index 0000000..c3980d3 --- /dev/null +++ b/dataproxy-server/src/test/java/org/secretflow/dataproxy/server/handler/MetricsChannelInboundHandlerTest.java @@ -0,0 +1,202 @@ +/* + * Copyright 2025 Ant Group Co., Ltd. + * + * 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 org.secretflow.dataproxy.server.handler; + +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelFutureListener; +import io.netty.channel.ChannelHandlerContext; +import io.netty.handler.codec.http.FullHttpRequest; +import io.netty.handler.codec.http.FullHttpResponse; +import io.netty.handler.codec.http.HttpHeaderNames; +import io.netty.handler.codec.http.HttpResponseStatus; +import io.netty.handler.codec.http.HttpVersion; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.nio.charset.StandardCharsets; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * @author yuexie + * @date 2025/6/6 14:12 + **/ +@ExtendWith(MockitoExtension.class) +public class MetricsChannelInboundHandlerTest { + + @Mock + private ChannelHandlerContext ctx; + + @Mock + private FullHttpRequest request; + + @Mock + private ChannelFuture channelFuture; + + @Captor + private ArgumentCaptor responseCaptor; + + // Test Case: Correctly respond to the /metrics request + @Test + void testHandleMetricsRequest() throws Exception { + // Given + when(request.uri()).thenReturn("/metrics"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + FullHttpResponse response = responseCaptor.getValue(); + + assertEquals(HttpResponseStatus.OK, response.status()); + assertEquals("text/plain; version=0.0.4;charset=utf-8", + response.headers().get(HttpHeaderNames.CONTENT_TYPE)); + assertTrue(response.content().readableBytes() > 0); + verify(channelFuture).addListener(ChannelFutureListener.CLOSE); + } + + // Test Case: Handling non/metrics path requests + @Test + void testHandleOtherPathRequest() throws Exception { + // Given + when(request.uri()).thenReturn("/other"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + FullHttpResponse response = responseCaptor.getValue(); + + assertEquals(HttpResponseStatus.NOT_FOUND, response.status()); + assertNull(response.headers().get(HttpHeaderNames.CONTENT_TYPE)); + assertEquals(0, response.content().readableBytes()); + verify(channelFuture).addListener(ChannelFutureListener.CLOSE); + } + + // Test Case: Processing /metrics requests with parameters + @Test + void testHandleMetricsWithQueryParams() throws Exception { + // Given + when(request.uri()).thenReturn("/metrics?debug=true"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + assertEquals(HttpResponseStatus.NOT_FOUND, responseCaptor.getValue().status()); + } + + // Test Case: Verify the encoding of the response content + @Test + void testResponseContentEncoding() throws Exception { + // Given + when(request.uri()).thenReturn("/metrics"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + String content = responseCaptor.getValue().content() + .toString(StandardCharsets.UTF_8); + assertTrue(content.contains("jvm_classes_loaded")); // 验证基础指标存在 + } + + // Test Case: Handle root path requests + @Test + void testHandleRootPathRequest() throws Exception { + // Given + when(request.uri()).thenReturn("/"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + assertEquals(HttpResponseStatus.NOT_FOUND, responseCaptor.getValue().status()); + } + + // Test Case: Verify the HTTP version + @Test + void testHttpProtocolVersion() throws Exception { + // Given + when(request.uri()).thenReturn("/metrics"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + assertEquals(HttpVersion.HTTP_1_1, responseCaptor.getValue().protocolVersion()); + } + + // Test Case: Handling all-caps URI requests + @Test + void testHandleUpperCaseUri() throws Exception { + // Given + when(request.uri()).thenReturn("/METRICS"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(ctx).writeAndFlush(responseCaptor.capture()); + assertEquals(HttpResponseStatus.NOT_FOUND, responseCaptor.getValue().status()); + } + + // Test Case: Verify the response to close the connection + @Test + void testResponseCloseConnection() throws Exception { + // Given + when(request.uri()).thenReturn("/metrics"); + when(ctx.writeAndFlush(any(FullHttpResponse.class))).thenReturn(channelFuture); + MetricsChannelInboundHandler handler = new MetricsChannelInboundHandler(); + + // When + handler.channelRead0(ctx, request); + + // Then + verify(channelFuture).addListener(ChannelFutureListener.CLOSE); + } + +} diff --git a/pom.xml b/pom.xml index f93d5c4..1dfdfc2 100644 --- a/pom.xml +++ b/pom.xml @@ -10,6 +10,7 @@ dataproxy-plugins dataproxy-core dataproxy-integration-tests + dataproxy-metrics org.secretflow @@ -43,7 +44,7 @@ 1.9.0 - 1.18.34 + 1.18.38 18.0.0 4.28.3 @@ -52,7 +53,7 @@ 2.11.0 2.36.0 3.0.2 - 4.1.115.Final + 4.2.4.Final 1.3.2 3.1.0 @@ -79,6 +80,8 @@ 5.2.0 2.1.7 5.11.4 + 1.14.5 + @@ -152,6 +155,15 @@ ${jakarta.validation-api.version} + + + io.micrometer + micrometer-bom + ${micrometer.version} + pom + import + + org.slf4j @@ -359,6 +371,11 @@ dataproxy-plugin-odps ${project.version} + + org.secretflow + dataproxy-metrics + ${project.version} + org.apache.hive hive-jdbc @@ -584,4 +601,4 @@ - \ No newline at end of file + diff --git a/scripts/build_image.sh b/scripts/build_image.sh index 030cadf..5e8ca6a 100755 --- a/scripts/build_image.sh +++ b/scripts/build_image.sh @@ -40,7 +40,7 @@ if [ "$BUILDER_EXISTS" -eq 0 ]; then docker buildx use dataproxy_image_buildx else echo "creating new buildx builder: dataproxy_image_buildx" - docker buildx create --name dataproxy_image_buildx --use + docker buildx create --name dataproxy_image_buildx --use --platform linux/arm64,linux/amd64 fi if [[ "$github_flag" == "true" ]]; then