From f9f64ea6c35b478286a737366eea5ac1192b37b5 Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Tue, 15 Sep 2026 16:47:33 +0800 Subject: [PATCH 1/6] [fix](ai) keep analyzed resource across same-name replacement --- .../org/apache/doris/catalog/Resource.java | 4 +- .../doris/nereids/StatementContext.java | 23 ++++++-- .../expressions/functions/agg/AIAgg.java | 2 +- .../expressions/functions/ai/AIFunction.java | 2 +- .../trees/expressions/functions/ai/Embed.java | 2 +- .../java/org/apache/doris/qe/Coordinator.java | 10 +--- .../doris/qe/runtime/ThriftPlansBuilder.java | 11 +--- .../qe/runtime/ThriftPlansBuilderTest.java | 56 +++++++++++++++++++ 8 files changed, 81 insertions(+), 29 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java index 92ec990f557e43..4417235e1d5208 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/Resource.java @@ -367,10 +367,10 @@ private void notifyUpdate(Map properties) { public void applyDefaultProperties() {} - public static void registerUsedAIResourceName(String resourceName) { + public static void registerUsedAIResource(AIResource resource) { ConnectContext ctx = ConnectContext.get(); if (ctx != null && ctx.getStatementContext() != null) { - ctx.getStatementContext().registerUsedAIResourceName(resourceName); + ctx.getStatementContext().registerUsedAIResource(resource.getName(), resource.toThrift()); } } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index 64670a7398b41a..593754acb8a4b3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -72,6 +72,7 @@ import org.apache.doris.qe.cache.CacheAnalyzer; import org.apache.doris.statistics.model.Statistics; import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TAIResource; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Stopwatch; @@ -354,7 +355,7 @@ public enum TableFrom { private boolean queryStatsRecorded = false; private final Set mustInlineCTE = new HashSet<>(); - private final Set usedAIResourceNames = new LinkedHashSet<>(); + private final Map usedAIResources = new LinkedHashMap<>(); private final Set excludedTriggerTables = new HashSet<>(); private final Map lowerCaseTableNamesCache = Maps.newHashMap(); @@ -591,15 +592,27 @@ public boolean isIvmMTMVRewrite() { return getIvmRewriteContext().isPresent(); } - public Set getUsedAIResourceNames() { - return Collections.unmodifiableSet(usedAIResourceNames); + public synchronized Set getUsedAIResourceNames() { + return Collections.unmodifiableSet(new LinkedHashSet<>(usedAIResources.keySet())); } - public void registerUsedAIResourceName(String resourceName) { + public synchronized Map getUsedAIResources() { + Map snapshots = new LinkedHashMap<>(); + usedAIResources.forEach((name, resource) -> snapshots.put(name, resource.deepCopy())); + return Collections.unmodifiableMap(snapshots); + } + + /** + * Retain the first validated configuration for an AI resource used by this statement. + */ + public synchronized void registerUsedAIResource(String resourceName, TAIResource resource) { if (Strings.isNullOrEmpty(resourceName)) { throw new AnalysisException("AI resource name can not be empty"); } - usedAIResourceNames.add(resourceName); + if (resource == null) { + throw new AnalysisException("AI resource snapshot can not be null"); + } + usedAIResources.putIfAbsent(resourceName, resource.deepCopy()); } /** diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java index 4785718b9faee8..8f162371e16fb1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java @@ -90,7 +90,7 @@ public void checkLegalityAfterRewrite() { if (!(resource instanceof AIResource)) { throw new AnalysisException("AI resource '" + resourceName + "' does not exist"); } - Resource.registerUsedAIResourceName(resourceName); + Resource.registerUsedAIResource((AIResource) resource); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java index 7399f2a348f68f..681ee7a60975cf 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java @@ -61,7 +61,7 @@ public void checkLegalityAfterRewrite() { if (!(resource instanceof AIResource)) { throw new AnalysisException("AI resource '" + resourceName + "' does not exist"); } - Resource.registerUsedAIResourceName(resourceName); + Resource.registerUsedAIResource((AIResource) resource); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java index 0082609b7fdb0d..1007c3ca197c54 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java @@ -124,7 +124,7 @@ private static void validateAIResource(String resourceName) { if (!(resource instanceof AIResource)) { throw new AnalysisException("AI resource '" + resourceName + "' does not exist"); } - Resource.registerUsedAIResourceName(resourceName); + Resource.registerUsedAIResource((AIResource) resource); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 0d3dec78561108..efdbe679e4ccae 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -21,10 +21,8 @@ import org.apache.doris.analysis.DescriptorToThriftConverter; import org.apache.doris.analysis.StorageBackend; import org.apache.doris.arrowflight.results.FlightSqlEndpointsLocation; -import org.apache.doris.catalog.AIResource; import org.apache.doris.catalog.Env; import org.apache.doris.catalog.FsBroker; -import org.apache.doris.catalog.Resource; import org.apache.doris.common.Config; import org.apache.doris.common.MarkedCountDownLatch; import org.apache.doris.common.Pair; @@ -1794,13 +1792,7 @@ private Map getNeededAiResources() { if (context == null || context.getStatementContext() == null) { return aiResourceMap; } - for (String resourceName : context.getStatementContext().getUsedAIResourceNames()) { - Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName); - if (!(resource instanceof AIResource)) { - throw new IllegalStateException("AI resource '" + resourceName + "' does not exist"); - } - aiResourceMap.put(resourceName, ((AIResource) resource).toThrift()); - } + aiResourceMap.putAll(context.getStatementContext().getUsedAIResources()); return aiResourceMap; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java index 64813f1b7b58e9..a76e16fec2be2f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java @@ -19,9 +19,6 @@ import org.apache.doris.analysis.Expr; import org.apache.doris.analysis.ExprToThriftVisitor; -import org.apache.doris.catalog.AIResource; -import org.apache.doris.catalog.Env; -import org.apache.doris.catalog.Resource; import org.apache.doris.common.Config; import org.apache.doris.datasource.scan.FileQueryScanNode; import org.apache.doris.nereids.StatementContext; @@ -230,13 +227,7 @@ static Map collectAiResources(ConnectContext connectContext return aiResourceMap; } - for (String resourceName : statementContext.getUsedAIResourceNames()) { - Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName); - if (!(resource instanceof AIResource)) { - throw new IllegalStateException("AI resource '" + resourceName + "' does not exist"); - } - aiResourceMap.put(resourceName, ((AIResource) resource).toThrift()); - } + aiResourceMap.putAll(statementContext.getUsedAIResources()); return aiResourceMap; } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java index 62a83172f94194..dfddc74d9ccdce 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/ThriftPlansBuilderTest.java @@ -19,6 +19,12 @@ import org.apache.doris.analysis.TupleDescriptor; import org.apache.doris.analysis.TupleId; +import org.apache.doris.catalog.AIResource; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.ResourceMgr; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.trees.expressions.functions.agg.AIAgg; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; import org.apache.doris.nereids.trees.plans.distribute.worker.DistributedPlanWorker; import org.apache.doris.nereids.trees.plans.distribute.worker.job.AssignedJob; import org.apache.doris.nereids.trees.plans.distribute.worker.job.DefaultScanSource; @@ -28,18 +34,68 @@ import org.apache.doris.planner.RecursiveCteScanNode; import org.apache.doris.planner.ScanNode; import org.apache.doris.planner.SortNode; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.OriginStatement; +import org.apache.doris.thrift.TAIResource; import org.apache.doris.thrift.TRecCTETarget; import org.apache.doris.thrift.TUniqueId; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; import org.mockito.Mockito; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; public class ThriftPlansBuilderTest { + @Test + public void testCollectAiResourcesKeepsAnalyzedSnapshotAfterResourceReplacement() { + String resourceName = "ai_resource"; + TAIResource analyzedThrift = new TAIResource() + .setProviderType("OPENAI") + .setModelName("analyzed-model"); + TAIResource replacementThrift = new TAIResource() + .setProviderType("LOCAL") + .setModelName("replacement-model"); + AIResource analyzedResource = Mockito.mock(AIResource.class); + AIResource replacementResource = Mockito.mock(AIResource.class); + Mockito.when(analyzedResource.getName()).thenReturn(resourceName); + Mockito.when(analyzedResource.toThrift()).thenReturn(analyzedThrift); + Mockito.when(replacementResource.toThrift()).thenReturn(replacementThrift); + + ResourceMgr resourceMgr = Mockito.mock(ResourceMgr.class); + Mockito.when(resourceMgr.getResource(resourceName)) + .thenReturn(analyzedResource, replacementResource); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getResourceMgr()).thenReturn(resourceMgr); + + ConnectContext previousContext = ConnectContext.get(); + ConnectContext connectContext = new ConnectContext(); + StatementContext statementContext = new StatementContext( + connectContext, new OriginStatement("select 1", 0)); + connectContext.setStatementContext(statementContext); + connectContext.setThreadLocalInfo(); + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + + new AIAgg(new StringLiteral(resourceName), new StringLiteral("value"), + new StringLiteral("summarize")).checkLegalityAfterRewrite(); + Map resources = ThriftPlansBuilder.collectAiResources(connectContext); + + Assertions.assertEquals("OPENAI", resources.get(resourceName).getProviderType()); + Assertions.assertEquals("analyzed-model", resources.get(resourceName).getModelName()); + Mockito.verify(resourceMgr, Mockito.times(1)).getResource(resourceName); + } finally { + ConnectContext.remove(); + if (previousContext != null) { + previousContext.setThreadLocalInfo(); + } + } + } + @Test public void testSetRuntimePredicateForNonOlapScanNode() { ScanNode scanNode = Mockito.mock(ScanNode.class); From 2262a246b199fdb041c1bb7564ab42795afa717c Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Tue, 15 Sep 2026 16:50:25 +0800 Subject: [PATCH 2/6] [fix](ai) send initial request when max retries is zero --- .../exprs/aggregate/aggregate_function_ai_agg.h | 2 +- be/src/exprs/function/ai/ai_adapter.h | 9 +++++++++ be/src/exprs/function/ai/ai_functions.h | 3 ++- be/test/ai/ai_adapter_test.cpp | 16 ++++++++++++++++ .../suites/ai_p0/test_create_ai_resource.groovy | 5 +++-- 5 files changed, 31 insertions(+), 4 deletions(-) diff --git a/be/src/exprs/aggregate/aggregate_function_ai_agg.h b/be/src/exprs/aggregate/aggregate_function_ai_agg.h index d3da899fca8557..3d92aa44676fc0 100644 --- a/be/src/exprs/aggregate/aggregate_function_ai_agg.h +++ b/be/src/exprs/aggregate/aggregate_function_ai_agg.h @@ -170,7 +170,7 @@ class AggregateFunctionAIAggData { #endif return HttpClient::execute_with_retry( - _ai_config.max_retries, _ai_config.retry_delay_second, + ai_http_request_attempts(_ai_config.max_retries), _ai_config.retry_delay_second, [this, &request_body, &response](HttpClient* client) -> Status { return this->do_send_request(client, request_body, response); }); diff --git a/be/src/exprs/function/ai/ai_adapter.h b/be/src/exprs/function/ai/ai_adapter.h index b6376c6c004fde..7028ff75e10f4b 100644 --- a/be/src/exprs/function/ai/ai_adapter.h +++ b/be/src/exprs/function/ai/ai_adapter.h @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -39,6 +40,14 @@ namespace doris { +// AI max_retries excludes the initial request, while HttpClient expects total attempts. +inline int ai_http_request_attempts(int32_t max_retries) { + if (max_retries <= 0) { + return 1; + } + return max_retries == std::numeric_limits::max() ? max_retries : max_retries + 1; +} + struct AIResource { AIResource() = default; AIResource(const TAIResource& tai) diff --git a/be/src/exprs/function/ai/ai_functions.h b/be/src/exprs/function/ai/ai_functions.h index 2ed287b37ad88c..598d1d53c5331c 100644 --- a/be/src/exprs/function/ai/ai_functions.h +++ b/be/src/exprs/function/ai/ai_functions.h @@ -209,7 +209,8 @@ class AIFunction : public IFunction { Status send_request_to_llm(const std::string& request_body, std::string& response, const TAIResource& config, std::shared_ptr& adapter, FunctionContext* context) const { - return HttpClient::execute_with_retry(config.max_retries, config.retry_delay_second, + return HttpClient::execute_with_retry(ai_http_request_attempts(config.max_retries), + config.retry_delay_second, [this, &request_body, &response, &config, &adapter, context](HttpClient* client) -> Status { return this->do_send_request(client, request_body, diff --git a/be/test/ai/ai_adapter_test.cpp b/be/test/ai/ai_adapter_test.cpp index da40ef217dc94a..9f4e15a67ace82 100644 --- a/be/test/ai/ai_adapter_test.cpp +++ b/be/test/ai/ai_adapter_test.cpp @@ -41,6 +41,22 @@ class MockHttpClient : public HttpClient { std::string _content_type; }; +TEST(AI_ADAPTER_TEST, max_retries_is_in_addition_to_initial_request) { + auto count_attempts = [](int32_t max_retries) { + int attempts = 0; + Status status = HttpClient::execute_with_retry( + ai_http_request_attempts(max_retries), 0, [&attempts](HttpClient*) { + ++attempts; + return Status::InternalError("force retry"); + }); + EXPECT_FALSE(status.ok()); + return attempts; + }; + + EXPECT_EQ(count_attempts(0), 1); + EXPECT_EQ(count_attempts(2), 3); +} + TEST(AI_ADAPTER_TEST, local_adapter_request_chat_endpoint) { LocalAdapter adapter; TAIResource config; diff --git a/regression-test/suites/ai_p0/test_create_ai_resource.groovy b/regression-test/suites/ai_p0/test_create_ai_resource.groovy index d9e9522ab6b0c1..7777d0d10f6fde 100644 --- a/regression-test/suites/ai_p0/test_create_ai_resource.groovy +++ b/regression-test/suites/ai_p0/test_create_ai_resource.groovy @@ -85,12 +85,13 @@ suite("test_create_ai_resource") { 'ai.api_key' = 'sk-xxx', 'ai.temperature' = '0.7', 'ai.max_token' = '1024', - 'ai.max_retries' = '3', + 'ai.max_retries' = '0', 'ai.retry_delay_second' = '1', 'ai.validity_check' = 'false' );""" def res = sql """SHOW RESOURCES WHERE NAME = '${resourceName}'""" assertTrue(res.size() > 0) + assertTrue(res.any { row -> row[2] == 'ai.max_retries' && row[3] == '0' }) try_sql("""DROP RESOURCE '${resourceName}'""") -} \ No newline at end of file +} From bc626c5899be9ee6747ec37f422fb15533790732 Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Tue, 15 Sep 2026 16:56:16 +0800 Subject: [PATCH 3/6] [fix](ai) reject invalid numeric resource properties --- .../property/constants/AIProperties.java | 83 ++++++++++++++++--- .../apache/doris/catalog/AIResourceTest.java | 37 +++++++++ .../ai_p0/test_create_ai_resource.groovy | 13 +++ 3 files changed, 121 insertions(+), 12 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java index 6019aaef303268..f6ff400bcfb8ba 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java @@ -78,22 +78,81 @@ public static void requiredAIProperties(Map properties) throws D + properties.get(AIProperties.PROVIDER_TYPE)); } - // Check weather the 'temperature' is valid + // Check whether the numeric properties are valid String temp = properties.get(AIProperties.TEMPERATURE); - if (!Strings.isNullOrEmpty(temp) && !temp.equals("-1")) { - double tempVal = Double.parseDouble(temp); - if (!(tempVal >= 0 && tempVal <= 1)) { - throw new DdlException("Temperature must be a double between 0 and 1"); + if (properties.containsKey(AIProperties.TEMPERATURE)) { + double tempVal; + try { + tempVal = Double.parseDouble(temp); + } catch (NullPointerException | NumberFormatException e) { + throw new DdlException("Temperature must be -1 or a double between 0 and 1", e); + } + if (!Double.isFinite(tempVal) || (tempVal != -1 && !(tempVal >= 0 && tempVal <= 1))) { + throw new DdlException("Temperature must be -1 or a double between 0 and 1"); } } - // Check 'dimensions' - temp = properties.get(AIProperties.DIMENSIONS); - if (!Strings.isNullOrEmpty(temp) && temp.equals("-1")) { - int tempVal = Integer.parseInt(temp); - if (tempVal <= 0) { - throw new DdlException("Dimensions must be a positive integer"); - } + if (properties.containsKey(AIProperties.MAX_TOKEN)) { + checkPositiveLongOrDefault(properties.get(AIProperties.MAX_TOKEN), "Max token"); + } + if (properties.containsKey(AIProperties.MAX_RETRIES)) { + checkMaxRetries(properties.get(AIProperties.MAX_RETRIES)); + } + if (properties.containsKey(AIProperties.RETRY_DELAY_SECOND)) { + checkNonNegativeInteger(properties.get(AIProperties.RETRY_DELAY_SECOND), "Retry delay second"); + } + if (properties.containsKey(AIProperties.DIMENSIONS)) { + checkPositiveIntegerOrDefault(properties.get(AIProperties.DIMENSIONS), "Dimensions"); + } + } + + private static void checkPositiveLongOrDefault(String value, String propertyName) throws DdlException { + long parsedValue; + try { + parsedValue = Long.parseLong(value); + } catch (NumberFormatException e) { + throw new DdlException(propertyName + " must be a positive integer or -1", e); + } + if (parsedValue != -1 && parsedValue <= 0) { + throw new DdlException(propertyName + " must be a positive integer or -1"); + } + } + + private static void checkNonNegativeInteger(String value, String propertyName) throws DdlException { + int parsedValue; + try { + parsedValue = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new DdlException(propertyName + " must be a non-negative integer", e); + } + if (parsedValue < 0) { + throw new DdlException(propertyName + " must be a non-negative integer"); + } + } + + private static void checkMaxRetries(String value) throws DdlException { + int parsedValue; + try { + parsedValue = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new DdlException("Max retries must be a non-negative integer less than " + + Integer.MAX_VALUE, e); + } + if (parsedValue < 0 || parsedValue == Integer.MAX_VALUE) { + throw new DdlException("Max retries must be a non-negative integer less than " + + Integer.MAX_VALUE); + } + } + + private static void checkPositiveIntegerOrDefault(String value, String propertyName) throws DdlException { + int parsedValue; + try { + parsedValue = Integer.parseInt(value); + } catch (NumberFormatException e) { + throw new DdlException(propertyName + " must be a positive integer or -1", e); + } + if (parsedValue != -1 && parsedValue <= 0) { + throw new DdlException(propertyName + " must be a positive integer or -1"); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java index 25e4a76e7b80c5..2c9e040798c4a8 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java @@ -314,6 +314,43 @@ public void testModifyProperties() throws Exception { Assertions.assertEquals("0.9", aiResource.getProperty(AIProperties.TEMPERATURE)); } + @Test + public void testRejectInvalidNumericProperties() { + assertInvalidNumericProperty(AIProperties.TEMPERATURE, "", "Temperature"); + assertInvalidNumericProperty(AIProperties.TEMPERATURE, "not-a-double", "Temperature"); + assertInvalidNumericProperty(AIProperties.TEMPERATURE, "NaN", "Temperature"); + assertInvalidNumericProperty(AIProperties.TEMPERATURE, "Infinity", "Temperature"); + assertInvalidNumericProperty(AIProperties.TEMPERATURE, "-Infinity", "Temperature"); + assertInvalidNumericProperty(AIProperties.TEMPERATURE, "1.1", "Temperature"); + assertInvalidNumericProperty(AIProperties.MAX_TOKEN, "", "Max token"); + assertInvalidNumericProperty(AIProperties.MAX_TOKEN, "not-a-long", "Max token"); + assertInvalidNumericProperty(AIProperties.MAX_TOKEN, "9223372036854775808", "Max token"); + assertInvalidNumericProperty(AIProperties.MAX_TOKEN, "0", "Max token"); + assertInvalidNumericProperty(AIProperties.MAX_RETRIES, "", "Max retries"); + assertInvalidNumericProperty(AIProperties.MAX_RETRIES, "not-an-int", "Max retries"); + assertInvalidNumericProperty(AIProperties.MAX_RETRIES, "2147483648", "Max retries"); + assertInvalidNumericProperty(AIProperties.MAX_RETRIES, "2147483647", "Max retries"); + assertInvalidNumericProperty(AIProperties.MAX_RETRIES, "-1", "Max retries"); + assertInvalidNumericProperty(AIProperties.RETRY_DELAY_SECOND, "", "Retry delay second"); + assertInvalidNumericProperty(AIProperties.RETRY_DELAY_SECOND, "not-an-int", "Retry delay second"); + assertInvalidNumericProperty(AIProperties.RETRY_DELAY_SECOND, "2147483648", "Retry delay second"); + assertInvalidNumericProperty(AIProperties.RETRY_DELAY_SECOND, "-1", "Retry delay second"); + assertInvalidNumericProperty(AIProperties.DIMENSIONS, "", "Dimensions"); + assertInvalidNumericProperty(AIProperties.DIMENSIONS, "not-an-int", "Dimensions"); + assertInvalidNumericProperty(AIProperties.DIMENSIONS, "2147483648", "Dimensions"); + assertInvalidNumericProperty(AIProperties.DIMENSIONS, "0", "Dimensions"); + } + + private void assertInvalidNumericProperty(String property, String value, String expectedMessage) { + Map properties = new HashMap<>(aiProperties); + properties.put(property, value); + AIResource aiResource = new AIResource("invalid-numeric-resource"); + + DdlException exception = Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + Assertions.assertTrue(exception.getMessage().contains(expectedMessage)); + } + @Test public void testDifferentProviders() throws DdlException { // 1. OpenAI diff --git a/regression-test/suites/ai_p0/test_create_ai_resource.groovy b/regression-test/suites/ai_p0/test_create_ai_resource.groovy index 7777d0d10f6fde..95a16e4e60b0eb 100644 --- a/regression-test/suites/ai_p0/test_create_ai_resource.groovy +++ b/regression-test/suites/ai_p0/test_create_ai_resource.groovy @@ -76,6 +76,19 @@ suite("test_create_ai_resource") { exception "Missing [ai.api_key] in properties for provider: DEEPSEEK" } + test { + sql """CREATE RESOURCE "${resourceName}" + PROPERTIES( + 'type' = 'ai', + 'ai.provider_type' = 'deepseek', + 'ai.endpoint' = 'https://api.deepseek.com/chat/completions', + 'ai.model_name' = 'deepseek-chat', + 'ai.api_key' = 'sk-xxx', + 'ai.max_retries' = '-1' + );""" + exception "Max retries must be a non-negative integer" + } + sql """CREATE RESOURCE IF NOT EXISTS "${resourceName}" PROPERTIES( 'type' = 'ai', From de201284b9e3ebfa4985e05d2b4da428ef33ce92 Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Tue, 15 Sep 2026 17:05:20 +0800 Subject: [PATCH 4/6] [fix](ai) remove AI resource validity check property --- .../org/apache/doris/catalog/AIResource.java | 74 +++++++++---------- .../property/constants/AIProperties.java | 2 - .../apache/doris/catalog/AIResourceTest.java | 54 +++++++++++--- .../suites/ai_p0/test_ai_functions.groovy | 3 +- .../ai_p0/test_create_ai_resource.groovy | 7 +- .../test_ddl_ai_resource_auth.groovy | 6 +- 6 files changed, 84 insertions(+), 62 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java index 346e58285fd372..f4eddf01f33d5a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java @@ -27,9 +27,8 @@ import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.gson.annotations.SerializedName; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; +import java.io.IOException; import java.util.HashMap; import java.util.Map; @@ -54,7 +53,8 @@ */ public class AIResource extends Resource { - private static final Logger LOG = LogManager.getLogger(AIResource.class); + private static final String LEGACY_VALIDITY_CHECK = "ai.validity_check"; + @SerializedName(value = "properties") private Map properties; @SerializedName(value = "createdByRoot") @@ -80,60 +80,56 @@ void setCreatedByRoot(boolean createdByRoot) { @Override protected void setProperties(ImmutableMap newProperties) throws DdlException { Preconditions.checkState(newProperties != null); - this.properties = Maps.newHashMap(newProperties); - - AIProperties.requiredAIProperties(properties); + Map changedProperties = Maps.newHashMap(newProperties); + changedProperties.remove(LEGACY_VALIDITY_CHECK); - boolean needCheck = isNeedCheck(properties); - if (LOG.isDebugEnabled()) { - LOG.debug("AI resource need check validity: {}", needCheck); - } - - AIProperties.optionalAIProperties(this.properties); + AIProperties.requiredAIProperties(changedProperties); + AIProperties.optionalAIProperties(changedProperties); + this.properties = changedProperties; } public String getProperty(String propertyKey) { return properties.get(propertyKey); } - private boolean isNeedCheck(Map newProperties) { - boolean needCheck = !this.properties.containsKey(AIProperties.VALIDITY_CHECK) - || Boolean.parseBoolean(this.properties.get(AIProperties.VALIDITY_CHECK)); - - if (newProperties != null && newProperties.containsKey(AIProperties.VALIDITY_CHECK)) { - needCheck = Boolean.parseBoolean(newProperties.get(AIProperties.VALIDITY_CHECK)); - } - - if ("LOCAL".equalsIgnoreCase(this.properties.getOrDefault(AIProperties.PROVIDER_TYPE, ""))) { - needCheck = false; - } - return needCheck; - } - @Override - public void modifyProperties(Map properties) throws DdlException { - boolean needCheck = isNeedCheck(properties); - if (LOG.isDebugEnabled()) { - LOG.debug("AI resource need check validity: {}", needCheck); - } - - if (needCheck) { - Map changedProperties = new HashMap<>(this.properties); - changedProperties.putAll(properties); - AIProperties.requiredAIProperties(changedProperties); + public void modifyProperties(Map newProperties) throws DdlException { + Map changedProperties = new HashMap<>(this.properties); + changedProperties.remove(LEGACY_VALIDITY_CHECK); + for (Map.Entry kv : newProperties.entrySet()) { + if (LEGACY_VALIDITY_CHECK.equals(kv.getKey())) { + continue; + } + replaceIfEffectiveValue(changedProperties, kv.getKey(), kv.getValue()); + if (AIProperties.API_KEY.equals(kv.getKey())) { + changedProperties.put(kv.getKey(), kv.getValue()); + } } + AIProperties.requiredAIProperties(changedProperties); + AIProperties.optionalAIProperties(changedProperties); // modify properties writeLock(); - for (Map.Entry kv : properties.entrySet()) { + for (Map.Entry kv : newProperties.entrySet()) { + if (LEGACY_VALIDITY_CHECK.equals(kv.getKey())) { + continue; + } replaceIfEffectiveValue(this.properties, kv.getKey(), kv.getValue()); - if (kv.getKey().equals(AIProperties.API_KEY)) { + if (AIProperties.API_KEY.equals(kv.getKey())) { this.properties.put(kv.getKey(), kv.getValue()); } } ++version; writeUnlock(); - super.modifyProperties(properties); + super.modifyProperties(newProperties); + } + + @Override + public void gsonPostProcess() throws IOException { + super.gsonPostProcess(); + if (properties != null) { + properties.remove(LEGACY_VALIDITY_CHECK); + } } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java index f6ff400bcfb8ba..f65b05f2c58638 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java @@ -50,8 +50,6 @@ public class AIProperties extends BaseProperties { public static final String DEFAULT_ANTHROPIC_VERSION = "2023-06-01"; public static final String DEFAULT_DIMENSIONS = "-1"; - public static final String VALIDITY_CHECK = "ai.validity_check"; - public static final List REQUIRED_FIELDS = Arrays.asList(ENDPOINT, PROVIDER_TYPE, MODEL_NAME); public static final List PROVIDERS = Arrays.asList("OPENAI", "LOCAL", "GEMINI", "DEEPSEEK", "ANTHROPIC", diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java index 2c9e040798c4a8..5a62638a6e9277 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java @@ -86,7 +86,6 @@ public void setUp() { aiProperties.put("ai.provider_type", providerType); aiProperties.put("ai.api_key", apiKey); aiProperties.put("ai.model_name", modelName); - aiProperties.put("ai.validity_check", "false"); } @Test @@ -256,13 +255,14 @@ public void testSerialization() throws Exception { "ai.endpoint", endpoint, "ai.provider_type", providerType, "ai.api_key", apiKey, - "ai.model_name", modelName, - "ai.validity_check", "false" + "ai.model_name", modelName ); AIResource aiResource2 = new AIResource("ai_2"); aiResource2.setCreatedByRoot(true); aiResource2.setProperties(properties); - aiResource2.write(aiDos); + JsonObject legacyAIResourceJson = JsonParser.parseString(GsonUtils.GSON.toJson(aiResource2)).getAsJsonObject(); + legacyAIResourceJson.getAsJsonObject("properties").addProperty("ai.validity_check", "false"); + Text.writeString(aiDos, legacyAIResourceJson.toString()); aiDos.flush(); aiDos.close(); @@ -286,6 +286,7 @@ public void testSerialization() throws Exception { Assertions.assertEquals(rAiResource2.getProperty(AIProperties.MAX_RETRIES), AIProperties.DEFAULT_MAX_RETRIES); Assertions.assertEquals(rAiResource2.getProperty(AIProperties.RETRY_DELAY_SECOND), AIProperties.DEFAULT_RETRY_DELAY_SECOND); + Assertions.assertNull(rAiResource2.getProperty("ai.validity_check")); // 3. delete aiDis.close(); @@ -298,8 +299,7 @@ public void testModifyProperties() throws Exception { "ai.endpoint", endpoint, "ai.provider_type", providerType, "ai.api_key", apiKey, - "ai.model_name", modelName, - "ai.validity_check", "false" + "ai.model_name", modelName ); AIResource aiResource = new AIResource("t_ai_source"); aiResource.setProperties(properties); @@ -314,6 +314,44 @@ public void testModifyProperties() throws Exception { Assertions.assertEquals("0.9", aiResource.getProperty(AIProperties.TEMPERATURE)); } + @Test + public void testLegacyValidityCheckIsIgnored() throws DdlException { + Map properties = new HashMap<>(aiProperties); + properties.put("ai.validity_check", "false"); + AIResource aiResource = new AIResource("legacy-validity-check-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + Assertions.assertNull(aiResource.getProperty("ai.validity_check")); + + aiResource.modifyProperties(ImmutableMap.of("ai.validity_check", "true")); + Assertions.assertNull(aiResource.getProperty("ai.validity_check")); + } + + @Test + public void testEndpointIsNotValidatedOnCreateOrAlter() throws DdlException { + Map properties = new HashMap<>(aiProperties); + properties.put(AIProperties.ENDPOINT, "not-a-url"); + AIResource aiResource = new AIResource("unchecked-endpoint-resource"); + + Assertions.assertDoesNotThrow(() -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + Assertions.assertDoesNotThrow(() -> aiResource.modifyProperties( + ImmutableMap.of(AIProperties.ENDPOINT, "still-not-a-url"))); + } + + @Test + public void testModifyProviderFromLocalRequiresApiKey() throws DdlException { + AIResource aiResource = new AIResource("local-resource-without-api-key"); + aiResource.setProperties(ImmutableMap.of( + AIProperties.ENDPOINT, "http://localhost:8000/v1/chat/completions", + AIProperties.PROVIDER_TYPE, "local", + AIProperties.MODEL_NAME, "local-model")); + + DdlException exception = Assertions.assertThrows(DdlException.class, () -> + aiResource.modifyProperties(ImmutableMap.of(AIProperties.PROVIDER_TYPE, "qwen"))); + + Assertions.assertTrue(exception.getMessage().contains("Missing [ai.api_key]")); + Assertions.assertEquals("LOCAL", aiResource.getProperty(AIProperties.PROVIDER_TYPE)); + } + @Test public void testRejectInvalidNumericProperties() { assertInvalidNumericProperty(AIProperties.TEMPERATURE, "", "Temperature"); @@ -359,7 +397,6 @@ public void testDifferentProviders() throws DdlException { openaiProps.put("ai.provider_type", "openai"); openaiProps.put("ai.api_key", "openai-key"); openaiProps.put("ai.model_name", "gpt-4"); - openaiProps.put("ai.validity_check", "false"); AIResource openaiResource = new AIResource("openai-resource"); openaiResource.setProperties(ImmutableMap.copyOf(openaiProps)); @@ -370,7 +407,6 @@ public void testDifferentProviders() throws DdlException { geminiProps.put("ai.provider_type", "gemini"); geminiProps.put("ai.api_key", "gemini-api-key"); geminiProps.put("ai.model_name", "gemini-pro"); - geminiProps.put("ai.validity_check", "false"); AIResource geminiResource = new AIResource("gemini-resource"); geminiResource.setProperties(ImmutableMap.copyOf(geminiProps)); @@ -382,7 +418,6 @@ public void testDifferentProviders() throws DdlException { anthropicProps.put("ai.api_key", "anthropic-api-key"); anthropicProps.put("ai.model_name", "claude-3-opus"); anthropicProps.put("ai.anthropic_version", "2023-06-01"); - anthropicProps.put("ai.validity_check", "false"); AIResource anthropicResource = new AIResource("anthropic-resource"); anthropicResource.setProperties(ImmutableMap.copyOf(anthropicProps)); @@ -393,7 +428,6 @@ public void testDifferentProviders() throws DdlException { localProps.put("ai.provider_type", "local"); localProps.put("ai.api_key", "local-key"); localProps.put("ai.model_name", "local-model"); - localProps.put("ai.validity_check", "false"); AIResource localResource = new AIResource("local-resource"); localResource.setProperties(ImmutableMap.copyOf(localProps)); diff --git a/regression-test/suites/ai_p0/test_ai_functions.groovy b/regression-test/suites/ai_p0/test_ai_functions.groovy index 2404573e398e66..ec509ba35d6dd0 100644 --- a/regression-test/suites/ai_p0/test_ai_functions.groovy +++ b/regression-test/suites/ai_p0/test_ai_functions.groovy @@ -37,8 +37,7 @@ suite("test_ai_functions") { 'ai.temperature' = '0.7', 'ai.max_token' = '1024', 'ai.max_retries' = '2', - 'ai.retry_delay_second' = '3', - 'ai.validity_check' = 'false' + 'ai.retry_delay_second' = '3' );""" def res = sql """SHOW RESOURCES WHERE NAME = '${resourceName}'""" diff --git a/regression-test/suites/ai_p0/test_create_ai_resource.groovy b/regression-test/suites/ai_p0/test_create_ai_resource.groovy index 95a16e4e60b0eb..390cddc18331c5 100644 --- a/regression-test/suites/ai_p0/test_create_ai_resource.groovy +++ b/regression-test/suites/ai_p0/test_create_ai_resource.groovy @@ -25,9 +25,6 @@ suite("test_create_ai_resource") { try_sql("""DROP RESOURCE '${resourceName}'""") - //If 'ai.validity_check'='false' is not set, - // ai resource availability must be checked when creating the resource. - // missing end_point test { sql """CREATE RESOURCE IF NOT EXISTS "${resourceName}" @@ -99,12 +96,12 @@ suite("test_create_ai_resource") { 'ai.temperature' = '0.7', 'ai.max_token' = '1024', 'ai.max_retries' = '0', - 'ai.retry_delay_second' = '1', - 'ai.validity_check' = 'false' + 'ai.retry_delay_second' = '1' );""" def res = sql """SHOW RESOURCES WHERE NAME = '${resourceName}'""" assertTrue(res.size() > 0) assertTrue(res.any { row -> row[2] == 'ai.max_retries' && row[3] == '0' }) + assertFalse(res.collect { row -> row[2] }.contains('ai.validity_check')) try_sql("""DROP RESOURCE '${resourceName}'""") } diff --git a/regression-test/suites/auth_call/test_ddl_ai_resource_auth.groovy b/regression-test/suites/auth_call/test_ddl_ai_resource_auth.groovy index 9046471382441c..1444eb1089e8f3 100644 --- a/regression-test/suites/auth_call/test_ddl_ai_resource_auth.groovy +++ b/regression-test/suites/auth_call/test_ddl_ai_resource_auth.groovy @@ -78,8 +78,7 @@ suite("test_ddl_ai_resource_auth","p0,auth_call") { 'ai.temperature' = '0.7', 'ai.max_token' = '1024', 'ai.max_retries' = '3', - 'ai.retry_delay_second' = '1', - 'ai.validity_check' = 'false' + 'ai.retry_delay_second' = '1' );""" def res = sql """SHOW RESOURCES WHERE NAME = '${resourceName}'""" assertTrue(res.size() > 0) @@ -99,8 +98,7 @@ suite("test_ddl_ai_resource_auth","p0,auth_call") { 'ai.temperature' = '0.7', 'ai.max_token' = '1024', 'ai.max_retries' = '3', - 'ai.retry_delay_second' = '1', - 'ai.validity_check' = 'false' + 'ai.retry_delay_second' = '1' );""" connect(user, "${pwd}", context.config.jdbcUrl) { test { From b944eff0299a2b32974a4bcfbbbcf9d590e25878 Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Tue, 15 Sep 2026 17:14:01 +0800 Subject: [PATCH 5/6] [fix](ai) persist normalized provider changes --- .../org/apache/doris/catalog/AIResource.java | 10 +------ .../property/constants/AIProperties.java | 6 +++-- .../apache/doris/catalog/AIResourceTest.java | 27 +++++++++++++++++++ 3 files changed, 32 insertions(+), 11 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java index f4eddf01f33d5a..11aa444e95e94b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java @@ -110,15 +110,7 @@ public void modifyProperties(Map newProperties) throws DdlExcept // modify properties writeLock(); - for (Map.Entry kv : newProperties.entrySet()) { - if (LEGACY_VALIDITY_CHECK.equals(kv.getKey())) { - continue; - } - replaceIfEffectiveValue(this.properties, kv.getKey(), kv.getValue()); - if (AIProperties.API_KEY.equals(kv.getKey())) { - this.properties.put(kv.getKey(), kv.getValue()); - } - } + this.properties = changedProperties; ++version; writeUnlock(); super.modifyProperties(newProperties); diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java index f65b05f2c58638..71047f5fcf2456 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java @@ -23,6 +23,7 @@ import java.util.Arrays; import java.util.List; +import java.util.Locale; import java.util.Map; public class AIProperties extends BaseProperties { @@ -64,8 +65,9 @@ public static void requiredAIProperties(Map properties) throws D } // Check the provider is valid - properties.put(PROVIDER_TYPE, properties.get(PROVIDER_TYPE).toUpperCase()); - if (PROVIDERS.stream().noneMatch(s -> s.equals(properties.get(PROVIDER_TYPE).toUpperCase()))) { + String provider = properties.get(PROVIDER_TYPE).toUpperCase(Locale.ROOT); + properties.put(PROVIDER_TYPE, provider); + if (!PROVIDERS.contains(provider)) { throw new DdlException("Provider must be one of " + PROVIDERS); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java index 5a62638a6e9277..0145a0082f0a9d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java @@ -32,6 +32,7 @@ import org.apache.doris.persist.EditLog; import org.apache.doris.persist.gson.GsonUtils; import org.apache.doris.qe.ConnectContext; +import org.apache.doris.thrift.TAIResource; import com.google.common.collect.ImmutableMap; import com.google.gson.JsonObject; @@ -44,6 +45,8 @@ import org.mockito.MockedStatic; import org.mockito.Mockito; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.DataInputStream; import java.io.DataOutputStream; import java.nio.file.Files; @@ -352,6 +355,30 @@ public void testModifyProviderFromLocalRequiresApiKey() throws DdlException { Assertions.assertEquals("LOCAL", aiResource.getProperty(AIProperties.PROVIDER_TYPE)); } + @Test + public void testModifyProviderStoresNormalizedValue() throws Exception { + AIResource aiResource = new AIResource("normalize-provider-resource"); + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + + aiResource.modifyProperties(ImmutableMap.of(AIProperties.PROVIDER_TYPE, "qwen")); + + Assertions.assertEquals("QWEN", aiResource.getProperty(AIProperties.PROVIDER_TYPE)); + TAIResource thriftResource = aiResource.toThrift(); + Assertions.assertEquals("QWEN", thriftResource.getProviderType()); + + ByteArrayOutputStream output = new ByteArrayOutputStream(); + try (DataOutputStream dataOutput = new DataOutputStream(output)) { + aiResource.write(dataOutput); + } + AIResource replayedResource; + try (DataInputStream dataInput = new DataInputStream( + new ByteArrayInputStream(output.toByteArray()))) { + replayedResource = (AIResource) Resource.read(dataInput); + } + Assertions.assertEquals("QWEN", replayedResource.getProperty(AIProperties.PROVIDER_TYPE)); + Assertions.assertEquals("QWEN", replayedResource.toThrift().getProviderType()); + } + @Test public void testRejectInvalidNumericProperties() { assertInvalidNumericProperty(AIProperties.TEMPERATURE, "", "Temperature"); From 29eba83824e756e3dba5913c4c6dd26f7baa5012 Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Tue, 15 Sep 2026 17:23:15 +0800 Subject: [PATCH 6/6] [fix](ai) prevent mixed resource state during concurrent ALTER --- .../org/apache/doris/catalog/AIResource.java | 82 +++++++---- .../apache/doris/catalog/AIResourceTest.java | 131 ++++++++++++++++++ .../ai_p0/test_create_ai_resource.groovy | 15 ++ 3 files changed, 198 insertions(+), 30 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java index 11aa444e95e94b..3c48b4113f7ff6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java @@ -56,7 +56,7 @@ public class AIResource extends Resource { private static final String LEGACY_VALIDITY_CHECK = "ai.validity_check"; @SerializedName(value = "properties") - private Map properties; + private volatile Map properties; @SerializedName(value = "createdByRoot") private boolean createdByRoot; @@ -66,7 +66,7 @@ public AIResource() { public AIResource(String name) { super(name, ResourceType.AI); - properties = Maps.newHashMap(); + properties = ImmutableMap.of(); } public boolean isCreatedByRoot() { @@ -85,34 +85,41 @@ protected void setProperties(ImmutableMap newProperties) throws AIProperties.requiredAIProperties(changedProperties); AIProperties.optionalAIProperties(changedProperties); - this.properties = changedProperties; + this.properties = ImmutableMap.copyOf(changedProperties); } public String getProperty(String propertyKey) { - return properties.get(propertyKey); + readLock(); + try { + return properties.get(propertyKey); + } finally { + readUnlock(); + } } @Override public void modifyProperties(Map newProperties) throws DdlException { - Map changedProperties = new HashMap<>(this.properties); - changedProperties.remove(LEGACY_VALIDITY_CHECK); - for (Map.Entry kv : newProperties.entrySet()) { - if (LEGACY_VALIDITY_CHECK.equals(kv.getKey())) { - continue; - } - replaceIfEffectiveValue(changedProperties, kv.getKey(), kv.getValue()); - if (AIProperties.API_KEY.equals(kv.getKey())) { - changedProperties.put(kv.getKey(), kv.getValue()); + writeLock(); + try { + Map changedProperties = new HashMap<>(this.properties); + changedProperties.remove(LEGACY_VALIDITY_CHECK); + for (Map.Entry kv : newProperties.entrySet()) { + if (LEGACY_VALIDITY_CHECK.equals(kv.getKey())) { + continue; + } + replaceIfEffectiveValue(changedProperties, kv.getKey(), kv.getValue()); + if (AIProperties.API_KEY.equals(kv.getKey())) { + changedProperties.put(kv.getKey(), kv.getValue()); + } } - } - AIProperties.requiredAIProperties(changedProperties); - AIProperties.optionalAIProperties(changedProperties); + AIProperties.requiredAIProperties(changedProperties); + AIProperties.optionalAIProperties(changedProperties); - // modify properties - writeLock(); - this.properties = changedProperties; - ++version; - writeUnlock(); + this.properties = ImmutableMap.copyOf(changedProperties); + ++version; + } finally { + writeUnlock(); + } super.modifyProperties(newProperties); } @@ -120,13 +127,15 @@ public void modifyProperties(Map newProperties) throws DdlExcept public void gsonPostProcess() throws IOException { super.gsonPostProcess(); if (properties != null) { - properties.remove(LEGACY_VALIDITY_CHECK); + Map loadedProperties = Maps.newHashMap(properties); + loadedProperties.remove(LEGACY_VALIDITY_CHECK); + properties = ImmutableMap.copyOf(loadedProperties); } } @Override public Map getCopiedProperties() { - return Maps.newHashMap(properties); + return getPropertiesSnapshot(); } @Override @@ -134,18 +143,22 @@ protected void getProcNodeData(BaseProcResult result) { String lowerCaseType = type.name().toLowerCase(); result.addRow(Lists.newArrayList(name, lowerCaseType, "id", String.valueOf(id))); readLock(); - result.addRow(Lists.newArrayList(name, lowerCaseType, "version", String.valueOf(version))); - for (Map.Entry entry : properties.entrySet()) { - if (entry.getKey().equals(AIProperties.API_KEY)) { - result.addRow(Lists.newArrayList(name, lowerCaseType, entry.getKey(), "******")); - } else { - result.addRow(Lists.newArrayList(name, lowerCaseType, entry.getKey(), entry.getValue())); + try { + result.addRow(Lists.newArrayList(name, lowerCaseType, "version", String.valueOf(version))); + for (Map.Entry entry : properties.entrySet()) { + if (entry.getKey().equals(AIProperties.API_KEY)) { + result.addRow(Lists.newArrayList(name, lowerCaseType, entry.getKey(), "******")); + } else { + result.addRow(Lists.newArrayList(name, lowerCaseType, entry.getKey(), entry.getValue())); + } } + } finally { + readUnlock(); } - readUnlock(); } public TAIResource toThrift() throws NumberFormatException { + Map properties = getPropertiesSnapshot(); TAIResource tAIResource = new TAIResource(); tAIResource.setProviderType(properties.get(AIProperties.PROVIDER_TYPE)); tAIResource.setEndpoint(properties.get(AIProperties.ENDPOINT)); @@ -186,4 +199,13 @@ public TAIResource toThrift() throws NumberFormatException { return tAIResource; } + + private Map getPropertiesSnapshot() { + readLock(); + try { + return Maps.newHashMap(properties); + } finally { + readUnlock(); + } + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java index 0145a0082f0a9d..ea76f10b8c4ba7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java @@ -54,6 +54,12 @@ import java.nio.file.Paths; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; public class AIResourceTest { private static final Logger LOG = LogManager.getLogger(AIResourceTest.class); @@ -379,6 +385,131 @@ public void testModifyProviderStoresNormalizedValue() throws Exception { Assertions.assertEquals("QWEN", replayedResource.toThrift().getProviderType()); } + @Test + public void testModifyPropertiesWithDefaultDimensions() throws DdlException { + AIResource aiResource = new AIResource("default-dimensions-resource"); + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + + Assertions.assertEquals(AIProperties.DEFAULT_DIMENSIONS, + aiResource.getProperty(AIProperties.DIMENSIONS)); + Assertions.assertDoesNotThrow(() -> aiResource.modifyProperties( + ImmutableMap.of(AIProperties.PROVIDER_TYPE, "QWEN"))); + + DdlException exception = Assertions.assertThrows(DdlException.class, () -> + aiResource.modifyProperties(ImmutableMap.of(AIProperties.DIMENSIONS, "0"))); + Assertions.assertTrue(exception.getMessage().contains("Dimensions must be a positive integer")); + } + + @Test + public void testInvalidNumericAlterIsAtomic() throws DdlException { + AIResource aiResource = new AIResource("invalid-numeric-alter-resource"); + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + long originalVersion = aiResource.version; + + DdlException exception = Assertions.assertThrows(DdlException.class, () -> + aiResource.modifyProperties(ImmutableMap.of(AIProperties.MAX_RETRIES, "not-an-int"))); + + Assertions.assertTrue(exception.getMessage().contains("Max retries")); + Assertions.assertEquals(AIProperties.DEFAULT_MAX_RETRIES, + aiResource.getProperty(AIProperties.MAX_RETRIES)); + Assertions.assertEquals(originalVersion, aiResource.version); + } + + @Test + public void testReadersWaitForConcurrentAlter() throws Exception { + CountDownLatch readLockAttempts = new CountDownLatch(3); + AIResource aiResource = new AIResource("concurrent-reader-resource") { + @Override + public void readLock() { + readLockAttempts.countDown(); + super.readLock(); + } + }; + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + ExecutorService executor = Executors.newFixedThreadPool(3); + + try { + aiResource.writeLock(); + Future property = executor.submit(() -> aiResource.getProperty(AIProperties.ENDPOINT)); + Future> copiedProperties = executor.submit(aiResource::getCopiedProperties); + Future thrift = executor.submit(aiResource::toThrift); + try { + Assertions.assertTrue(readLockAttempts.await(5, TimeUnit.SECONDS)); + Assertions.assertFalse(property.isDone()); + Assertions.assertFalse(copiedProperties.isDone()); + Assertions.assertFalse(thrift.isDone()); + } finally { + aiResource.writeUnlock(); + } + + Assertions.assertEquals(endpoint, property.get(5, TimeUnit.SECONDS)); + Assertions.assertEquals(endpoint, copiedProperties.get(5, TimeUnit.SECONDS).get(AIProperties.ENDPOINT)); + Assertions.assertEquals(endpoint, thrift.get(5, TimeUnit.SECONDS).getEndpoint()); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void testConcurrentAltersMergeAgainstLatestSnapshot() throws Exception { + CountDownLatch firstWriteLockAttempt = new CountDownLatch(1); + CountDownLatch secondWriteLockAttempt = new CountDownLatch(1); + CountDownLatch firstAlterFinished = new CountDownLatch(1); + AtomicInteger writeLockAttempts = new AtomicInteger(); + AIResource aiResource = new AIResource("concurrent-alter-resource") { + @Override + public void writeLock() { + int attempt = writeLockAttempts.incrementAndGet(); + if (attempt == 1) { + firstWriteLockAttempt.countDown(); + awaitLatch(secondWriteLockAttempt); + } else if (attempt == 2) { + secondWriteLockAttempt.countDown(); + awaitLatch(firstAlterFinished); + } + super.writeLock(); + } + + @Override + public void writeUnlock() { + super.writeUnlock(); + firstAlterFinished.countDown(); + } + }; + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + ExecutorService executor = Executors.newFixedThreadPool(2); + + try { + Future endpointAlter = executor.submit(() -> { + aiResource.modifyProperties(ImmutableMap.of(AIProperties.ENDPOINT, "https://new-endpoint")); + return null; + }); + Assertions.assertTrue(firstWriteLockAttempt.await(5, TimeUnit.SECONDS)); + Future temperatureAlter = executor.submit(() -> { + aiResource.modifyProperties(ImmutableMap.of(AIProperties.TEMPERATURE, "0.8")); + return null; + }); + + endpointAlter.get(5, TimeUnit.SECONDS); + temperatureAlter.get(5, TimeUnit.SECONDS); + Assertions.assertEquals("https://new-endpoint", aiResource.getProperty(AIProperties.ENDPOINT)); + Assertions.assertEquals("0.8", aiResource.getProperty(AIProperties.TEMPERATURE)); + } finally { + executor.shutdownNow(); + } + } + + private static void awaitLatch(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) { + throw new AssertionError("Timed out waiting for concurrent test coordination"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while coordinating concurrent test", e); + } + } + @Test public void testRejectInvalidNumericProperties() { assertInvalidNumericProperty(AIProperties.TEMPERATURE, "", "Temperature"); diff --git a/regression-test/suites/ai_p0/test_create_ai_resource.groovy b/regression-test/suites/ai_p0/test_create_ai_resource.groovy index 390cddc18331c5..fb172b90ca8e16 100644 --- a/regression-test/suites/ai_p0/test_create_ai_resource.groovy +++ b/regression-test/suites/ai_p0/test_create_ai_resource.groovy @@ -102,6 +102,21 @@ suite("test_create_ai_resource") { assertTrue(res.size() > 0) assertTrue(res.any { row -> row[2] == 'ai.max_retries' && row[3] == '0' }) assertFalse(res.collect { row -> row[2] }.contains('ai.validity_check')) + def propertiesBeforeInvalidAlter = res.collectEntries { row -> [(row[2]): row[3]] } + + test { + sql """ALTER RESOURCE "${resourceName}" PROPERTIES ('ai.dimensions' = '0')""" + exception "Dimensions must be a positive integer or -1" + } + + def propertiesAfterInvalidAlter = (sql """SHOW RESOURCES WHERE NAME = '${resourceName}'""") + .collectEntries { row -> [(row[2]): row[3]] } + assertEquals(propertiesBeforeInvalidAlter, propertiesAfterInvalidAlter) + + sql """ALTER RESOURCE "${resourceName}" PROPERTIES ('ai.max_retries' = '2')""" + def propertiesAfterValidAlter = (sql """SHOW RESOURCES WHERE NAME = '${resourceName}'""") + .collectEntries { row -> [(row[2]): row[3]] } + assertEquals('2', propertiesAfterValidAlter['ai.max_retries']) try_sql("""DROP RESOURCE '${resourceName}'""") }