From cffbf5b5a677a702d66d814d9983016c60e086d5 Mon Sep 17 00:00:00 2001 From: Zihan Dai Date: Tue, 10 Nov 2026 01:37:08 +1100 Subject: [PATCH] Fix table-model UDFs on TAG/ATTRIBUTE predicates and SHOW/COUNT DEVICES A table scalar UDF creates its IoTDBLocal when the transformer is built, which needs a factory, a fragment instance id, a query id and a positive deadline. Schema fragments had no deadline, and the COUNT DEVICES filter got none of the four, so any UDF evaluated there failed. Give schema fragments the same deadline as data fragments, and pass the four values to the COUNT DEVICES filter. --- .../relational/it/db/it/udf/IoTDBLocalIT.java | 92 +++++++++++++++++++ .../fragment/FragmentInstanceContext.java | 2 + .../fragment/FragmentInstanceManager.java | 1 + .../DataNodeTableOperatorGenerator.java | 6 +- 4 files changed, 100 insertions(+), 1 deletion(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java index 4e918e8bbef99..783b004ada83a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBLocalIT.java @@ -20,9 +20,11 @@ package org.apache.iotdb.relational.it.db.it.udf; import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.itbase.category.TableClusterIT; import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; +import org.apache.iotdb.itbase.env.BaseEnv; import org.junit.After; import org.junit.AfterClass; @@ -32,10 +34,12 @@ import org.junit.runner.RunWith; import java.sql.Connection; +import java.sql.SQLException; import java.sql.Statement; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.junit.Assume.assumeTrue; @@ -65,6 +69,8 @@ public class IoTDBLocalIT { "CREATE TABLE secret_table (device_id STRING TAG, secret STRING FIELD)", "CREATE TABLE vehicle (device_id STRING TAG, s1 INT32 FIELD, s2 INT64 FIELD)", "CREATE TABLE probe_table (device_id STRING TAG, value DOUBLE FIELD)", + "CREATE TABLE device_attrs (device_id STRING TAG, alias STRING ATTRIBUTE, v INT32 FIELD)", + "CREATE TABLE nullable_attrs (device_id STRING TAG, alias STRING ATTRIBUTE, v INT32 FIELD)", "INSERT INTO device_info(time, device_id, device_name) VALUES (1, 'd1', '一号车间温度传感器'), (1, 'd2', '二号车间温度传感器')", "INSERT INTO device_limits(time, device_id, max_temp) VALUES (1, 'd1', 30.0), (1, 'd2', 35.0)", "INSERT INTO readings(time, device_id, temperature) VALUES (1000, 'd1', 25.5), (1001, 'd2', 32.0), (1002, 'd3', 20.0)", @@ -74,6 +80,9 @@ public class IoTDBLocalIT { "INSERT INTO vehicle(time, device_id, s1, s2) VALUES (2, 'd0', null, 2)", "INSERT INTO vehicle(time, device_id, s1, s2) VALUES (3, 'd0', 3, 3)", "INSERT INTO vehicle(time, device_id, s1) VALUES (5, 'd0', 4)", + "INSERT INTO device_attrs(time, device_id, alias, v) VALUES (1, 'x1', 'd1', 1), (2, 'x1', 'd1', 2), (1, 'x2', 'd2', 3)", + "INSERT INTO nullable_attrs(time, device_id, alias, v) VALUES (1, 'y1', 'd1', 1)", + "INSERT INTO nullable_attrs(time, device_id, v) VALUES (1, 'y2', 2)", "FLUSH", "CLEAR ATTRIBUTE CACHE", }; @@ -106,6 +115,58 @@ private static void executeAsRoot(String... sqls) { } } + // ── scalar UDF on TAG/ATTRIBUTE columns, evaluated in schema fragments ───── + + @Test + public void testLegacyScalarUdfOnDeviceColumnPredicates() { + SQLFunctionUtils.createUDF("contain_null", PKG + ".ContainNull"); + tableResultSetEqualTest( + "SELECT count(*) AS cnt FROM readings WHERE contain_null(device_id) = false", + new String[] {"cnt"}, + new String[] {"3,"}, + DATABASE_NAME); + tableResultSetEqualTest( + "SELECT count(*) AS cnt FROM device_attrs WHERE contain_null(alias) = false", + new String[] {"cnt"}, + new String[] {"3,"}, + DATABASE_NAME); + tableResultSetEqualTest( + "SHOW DEVICES FROM nullable_attrs WHERE contain_null(alias) = true", + new String[] {"device_id", "alias"}, + new String[] {"y2,null,"}, + DATABASE_NAME); + tableResultSetEqualTest( + "COUNT DEVICES FROM device_attrs WHERE contain_null(alias) = false", + new String[] {"count(devices)"}, + new String[] {"2,"}, + DATABASE_NAME); + } + + @Test + public void testIoTDBLocalScalarUdfOnDeviceColumnPredicates() { + SQLFunctionUtils.createUDF("device_name", IOTDB_LOCAL_PKG + ".DeviceNameFunction"); + tableResultSetEqualTest( + "SELECT device_id, temperature FROM readings WHERE device_name(device_id) = '一号车间温度传感器'", + new String[] {"device_id", "temperature"}, + new String[] {"d1,25.5,"}, + DATABASE_NAME); + tableResultSetEqualTest( + "SELECT device_id, v FROM device_attrs WHERE device_name(alias) = '二号车间温度传感器'", + new String[] {"device_id", "v"}, + new String[] {"x2,3,"}, + DATABASE_NAME); + tableResultSetEqualTest( + "SHOW DEVICES FROM device_attrs WHERE device_name(alias) = '一号车间温度传感器'", + new String[] {"device_id", "alias"}, + new String[] {"x1,d1,"}, + DATABASE_NAME); + tableResultSetEqualTest( + "COUNT DEVICES FROM readings WHERE device_name(device_id) = '未知设备'", + new String[] {"count(devices)"}, + new String[] {"1,"}, + DATABASE_NAME); + } + // ── compatibility: legacy UDF without IoTDBLocal parameter ───────────────── @Test @@ -418,6 +479,37 @@ public void testIoTDBLocalDeniedWithoutTablePermission() { dropLimitedUser(); } + @Test + public void testIoTDBLocalOnDeviceColumnPredicateDeniedWithoutTablePermission() { + setupLimitedUserWithTableGrants("readings"); + SQLFunctionUtils.createUDF("secret_query", IOTDB_LOCAL_PKG + ".SecretTableQueryFunction"); + assertAccessDeniedOnEachDataNode( + "SELECT device_id FROM readings WHERE secret_query(device_id) = 'top-secret'"); + assertAccessDeniedOnEachDataNode( + "SHOW DEVICES FROM readings WHERE secret_query(device_id) = 'top-secret'"); + assertAccessDeniedOnEachDataNode( + "COUNT DEVICES FROM readings WHERE secret_query(device_id) = 'top-secret'"); + dropLimitedUser(); + } + + // A failed schema fragment can reach the client as 803 or 305, depending on whether the query + // state or the local source handle reports it first, so the DataNodes are checked one by one. + private static void assertAccessDeniedOnEachDataNode(String sql) { + for (DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) { + try (Connection connection = + EnvFactory.getEnv() + .getConnection( + dataNode, LIMITED_USER, LIMITED_PASSWORD, BaseEnv.TABLE_SQL_DIALECT); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE_NAME); + statement.executeQuery(sql); + fail("No exception: " + sql); + } catch (SQLException e) { + assertTrue(e.getMessage(), e.getMessage().contains("Access Denied")); + } + } + } + private static void setupLimitedUserWithTableGrants(String... tables) { try (Connection connection = EnvFactory.getEnv().getTableConnection(); Statement statement = connection.createStatement()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index ee4ff7377dd3f..74a1cf529d05c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -208,10 +208,12 @@ public static FragmentInstanceContext createFragmentInstanceContext( FragmentInstanceId id, FragmentInstanceStateMachine stateMachine, SessionInfo sessionInfo, + long outerQueryDeadlineMs, boolean debug, boolean isVerbose) { FragmentInstanceContext instanceContext = new FragmentInstanceContext(id, stateMachine, sessionInfo, debug, isVerbose); + instanceContext.outerQueryDeadlineMs = outerQueryDeadlineMs; instanceContext.initialize(); instanceContext.start(); return instanceContext; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java index fa6a45e837496..f5fb70c94da39 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java @@ -322,6 +322,7 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( fragmentInstanceId, stateMachine, instance.getSessionInfo(), + DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), instance.isDebug(), instance.isVerbose()); }); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java index 01584fe230776..44ca68490690f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java @@ -1444,7 +1444,11 @@ public Operator visitTableDeviceQueryCount( 0, context.getTypeProvider(), metadata, - context.getMemoryReservationManager())), + context.getMemoryReservationManager(), + getFragmentInstanceId(context), + getQueryId(context), + getOuterQueryDeadlineMs(context), + getIoTDBLocalFactory(context))), columnSchemaList, database, table)