Support 'set sql_dialect' in Tree and Table Model

apache · Feb 14, 2025 · 4d2ac44 · 4d2ac44
1 parent 3580c34
commit 4d2ac44
Show file tree

Hide file tree

Showing 33 changed files with 582 additions and 11 deletions.
diff --git a/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestStatement.java b/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestStatement.java
@@ -31,6 +31,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.apache.iotdb.rpc.RpcUtils.isSetSqlDialect;
 import static org.apache.iotdb.rpc.RpcUtils.isUseDatabase;
 
 /** The implementation of {@link ClusterTestStatement} in cluster test. */
@@ -191,8 +192,8 @@ public void setCursorName(String name) {
   public boolean execute(String sql) throws SQLException {
     sql = sql.trim();
     boolean result = writeStatement.execute(sql);
-    // if use XXXX, sendRequest to all statements
-    if (isUseDatabase(sql)) {
+    // if 'use XXXX' or 'set sql_dialect', sendRequest to all statements
+    if (isUseDatabase(sql) || isSetSqlDialect(sql)) {
       for (Statement readStatement : readStatements) {
         boolean tmp = readStatement.execute(sql);
         result = result && tmp;

diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/auth/IoTDBAuthIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/auth/IoTDBAuthIT.java
@@ -47,6 +47,8 @@
 import java.util.List;
 import java.util.Set;
 
+import static org.apache.iotdb.db.it.utils.TestUtils.createUser;
+import static org.apache.iotdb.db.it.utils.TestUtils.resultSetEqualTest;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -1278,4 +1280,15 @@ public void testCreateRoleIdentifierName() throws SQLException {
     adminStmt.execute("create role tail");
     adminStmt.execute("create user tail 'password'");
   }
+
+  @Test
+  public void noNeedPrivilegeTest() {
+    createUser("tempuser", "temppw");
+    String[] expectedHeader = new String[] {"CurrentUser"};
+    String[] retArray =
+        new String[] {
+          "tempuser,",
+        };
+    resultSetEqualTest("show current_user", expectedHeader, retArray, "tempuser", "temppw");
+  }
 }
diff --git a/...on-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBNullIdQueryIT.java b/...on-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBNullIdQueryIT.java
@@ -34,9 +34,11 @@
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.sql.Types;
 
+import static org.apache.iotdb.db.it.utils.TestUtils.createUser;
 import static org.apache.iotdb.db.it.utils.TestUtils.defaultFormatDataTime;
 import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
@@ -502,4 +504,35 @@ public void showStatementTest() {
     expectedHeader = new String[] {"CurrentTimestamp"};
     tableResultSetFuzzyTest("show current_timestamp", expectedHeader, 1, DATABASE_NAME);
   }
+
+  @Test
+  public void setSqlDialectTest() throws SQLException {
+    createUser("tempuser", "temppw");
+
+    try (Connection userCon = EnvFactory.getEnv().getConnection("tempuser", "temppw");
+        Statement userStmt = userCon.createStatement()) {
+      assertCurrentSqlDialect(true, userStmt);
+
+      // set Tree to Table
+      userStmt.execute("set sql_dialect=table");
+      assertCurrentSqlDialect(false, userStmt);
+
+      // set Table to Tree
+      userStmt.execute("set sql_dialect=tree");
+      assertCurrentSqlDialect(true, userStmt);
+    }
+  }
+
+  public static void assertCurrentSqlDialect(boolean expectedTree, Statement statement)
+      throws SQLException {
+    ResultSet resultSet = statement.executeQuery("show current_sql_dialect");
+    ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+    assertEquals("CurrentSqlDialect", resultSetMetaData.getColumnName(1));
+    int count = 0;
+    while (resultSet.next()) {
+      assertEquals(expectedTree ? "TREE" : "TABLE", resultSet.getString(1));
+      count++;
+    }
+    assertEquals(1, count);
+  }
 }
diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/Constant.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/Constant.java
@@ -42,6 +42,9 @@ private Constant() {}
   public static final String STATISTICS_RESULT_LINES = "* Lines of result: %d";
   public static final String STATISTICS_PRC_INFO = "* Num of RPC: %d, avg cost: %d ms";
 
+  public static final String TREE = "tree";
+  public static final String TABLE = "table";
+
   // version number
   public enum Version {
     V_0_12,

diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnection.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnection.java
@@ -617,6 +617,10 @@ protected void changeDefaultDatabase(String database) {
     params.setDb(database);
   }
 
+  protected void changeDefaultSqlDialect(String sqlDialect) {
+    params.setSqlDialect(sqlDialect);
+  }
+
   public int getTimeFactor() {
     return timeFactor;
   }

diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnectionParams.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBConnectionParams.java
@@ -27,6 +27,8 @@
 import java.time.ZoneId;
 import java.util.Optional;
 
+import static org.apache.iotdb.jdbc.Constant.TREE;
+
 public class IoTDBConnectionParams {
 
   private String host = Config.IOTDB_DEFAULT_HOST;
@@ -50,7 +52,7 @@ public class IoTDBConnectionParams {
   private String trustStore;
   private String trustStorePwd;
 
-  private String sqlDialect = "tree";
+  private String sqlDialect = TREE;
 
   private String db;
 

diff --git a/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBStatement.java b/iotdb-client/jdbc/src/main/java/org/apache/iotdb/jdbc/IoTDBStatement.java
@@ -46,6 +46,9 @@
 import java.util.BitSet;
 import java.util.List;
 
+import static org.apache.iotdb.jdbc.Constant.TABLE;
+import static org.apache.iotdb.jdbc.Constant.TREE;
+
 public class IoTDBStatement implements Statement {
 
   private final IoTDBConnection connection;
@@ -326,6 +329,10 @@ private boolean executeSQL(String sql) throws TException, SQLException {
       connection.changeDefaultDatabase(execResp.getDatabase());
     }
 
+    if (execResp.isSetTableModel()) {
+      connection.changeDefaultSqlDialect(execResp.tableModel ? TABLE : TREE);
+    }
+
     if (execResp.isSetColumns()) {
       queryId = execResp.getQueryId();
       if (execResp.queryResult == null) {

diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/RpcUtils.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/RpcUtils.java
@@ -374,6 +374,20 @@ public static boolean isUseDatabase(String sql) {
     return sql.length() > 4 && "use ".equalsIgnoreCase(sql.substring(0, 4));
   }
 
+  public static boolean isSetSqlDialect(String sql) {
+    // check if startWith 'set '
+    if (sql.length() <= 15 || !"set ".equalsIgnoreCase(sql.substring(0, 4))) {
+      return false;
+    }
+
+    // check if the following content of sql is 'sql_dialect'
+    sql = sql.substring(4).trim();
+    if (sql.length() <= 11) {
+      return false;
+    }
+    return sql.substring(0, 11).equalsIgnoreCase("sql_dialect");
+  }
+
   public static long getMilliSecond(long time, int timeFactor) {
     return time / timeFactor * 1_000;
   }

diff --git a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/RpcUtilsTest.java b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/RpcUtilsTest.java
@@ -64,4 +64,14 @@ public void parseLongToDateWithPrecision() {
         "1970-01-01T07:59:59.999+08:00",
         RpcUtils.parseLongToDateWithPrecision(formatter, -1, zoneId, "ms"));
   }
+
+  @Test
+  public void testIsSetSqlDialect() {
+    Assert.assertTrue(RpcUtils.isSetSqlDialect("set sql_dialect=table"));
+    Assert.assertTrue(RpcUtils.isSetSqlDialect("set sql_dialect =table"));
+    Assert.assertTrue(RpcUtils.isSetSqlDialect("set sql_dialect  =table"));
+    Assert.assertTrue(RpcUtils.isSetSqlDialect("set  sql_dialect =table"));
+    Assert.assertFalse(RpcUtils.isSetSqlDialect("setsql_dialect =table"));
+    Assert.assertFalse(RpcUtils.isSetSqlDialect("set           sql_dia"));
+  }
 }
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java
@@ -192,7 +192,7 @@ public class Session implements ISession {
 
   protected long retryIntervalInMs = SessionConfig.RETRY_INTERVAL_IN_MS;
 
-  protected String sqlDialect = SessionConfig.SQL_DIALECT;
+  protected volatile String sqlDialect = SessionConfig.SQL_DIALECT;
 
   // may be null
   protected volatile String database;
@@ -215,6 +215,9 @@ public class Session implements ISession {
       "All values are null and this submission is ignored,deviceIds are [{}],times are [{}],measurements are [{}]";
   private static final String ALL_INSERT_DATA_IS_NULL = "All inserted data is null.";
 
+  protected static final String TABLE = "table";
+  protected static final String TREE = "tree";
+
   public Session(String host, int rpcPort) {
     this(
         host,
@@ -994,8 +997,10 @@ private SessionConnection getQuerySessionConnection() {
   public void executeNonQueryStatement(String sql)
       throws IoTDBConnectionException, StatementExecutionException {
     String previousDB = database;
+    String previousDialect = sqlDialect;
     defaultSessionConnection.executeNonQueryStatement(sql);
-    if (!Objects.equals(previousDB, database) && endPointToSessionConnection != null) {
+    if ((!Objects.equals(previousDB, database) || !Objects.equals(previousDialect, sqlDialect))
+        && endPointToSessionConnection != null) {
       Iterator<Map.Entry<TEndPoint, SessionConnection>> iterator =
           endPointToSessionConnection.entrySet().iterator();
       while (iterator.hasNext()) {
@@ -1005,7 +1010,7 @@ public void executeNonQueryStatement(String sql)
           try {
             sessionConnection.executeNonQueryStatement(sql);
           } catch (Throwable t) {
-            logger.warn("failed to change database for {}", entry.getKey());
+            logger.warn("failed to execute '{}' for {}", sql, entry.getKey());
             iterator.remove();
           }
         }
@@ -4156,6 +4161,14 @@ public String getDatabase() {
     return database;
   }
 
+  protected void changeSqlDialect(String sqlDialect) {
+    this.sqlDialect = sqlDialect;
+  }
+
+  public String getSqlDialect() {
+    return sqlDialect;
+  }
+
   public static class Builder extends AbstractSessionBuilder {
 
     public Builder host(String host) {

diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
@@ -85,6 +85,9 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
+import static org.apache.iotdb.session.Session.TABLE;
+import static org.apache.iotdb.session.Session.TREE;
+
 @SuppressWarnings("java:S2142")
 public class SessionConnection {
 
@@ -107,7 +110,7 @@ public class SessionConnection {
 
   private final long retryIntervalInMs;
 
-  private final String sqlDialect;
+  private String sqlDialect;
 
   private String database;
 
@@ -471,6 +474,13 @@ private TSStatus executeNonQueryStatementInternal(TSExecuteStatementReq request)
       session.changeDatabase(dbName);
       this.database = dbName;
     }
+    if (resp.isSetTableModel()) {
+      String sqlDialect = resp.tableModel ? TABLE : TREE;
+      if (!sqlDialect.equalsIgnoreCase(this.sqlDialect)) {
+        session.changeSqlDialect(sqlDialect);
+        this.sqlDialect = sqlDialect;
+      }
+    }
     return resp.status;
   }
 

diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java
@@ -27,6 +27,8 @@
 import java.util.Collections;
 import java.util.List;
 
+import static org.apache.iotdb.session.Session.TABLE;
+
 /**
  * A builder class for constructing instances of {@link ITableSession}.
  *
@@ -267,7 +269,7 @@ public ITableSession build() throws IoTDBConnectionException {
       this.nodeUrls =
           Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT);
     }
-    this.sqlDialect = "table";
+    this.sqlDialect = TABLE;
     Session newSession = new Session(this);
     newSession.open(enableCompression, connectionTimeoutInMs);
     return new TableSession(newSession);

diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java
@@ -59,6 +59,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 
+import static org.apache.iotdb.rpc.RpcUtils.isSetSqlDialect;
 import static org.apache.iotdb.rpc.RpcUtils.isUseDatabase;
 
 /**
@@ -3080,8 +3081,8 @@ public SessionDataSetWrapper executeQueryStatement(String sql, long timeoutInMs)
   public void executeNonQueryStatement(String sql)
       throws StatementExecutionException, IoTDBConnectionException {
 
-    // use XXX is forbidden in SessionPool.executeNonQueryStatement
-    if (isUseDatabase(sql)) {
+    // 'use XXX' and 'set sql_dialect' is forbidden in SessionPool.executeNonQueryStatement
+    if (isUseDatabase(sql) || isSetSqlDialect(sql)) {
       throw new IllegalArgumentException(
           String.format("SessionPool doesn't support executing %s directly", sql));
     }

diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionWrapper.java
@@ -121,6 +121,21 @@ public void close() throws IoTDBConnectionException {
           return;
         }
       }
+
+      if (!Objects.equals(session.getSqlDialect(), sessionPool.sqlDialect)) {
+        try {
+          session.executeNonQueryStatement("set sql_dialect=" + sessionPool.sqlDialect);
+        } catch (StatementExecutionException e) {
+          LOGGER.warn(
+              "Failed to change back sql_dialect by executing: set sql_dialect={}",
+              sessionPool.sqlDialect,
+              e);
+          session.close();
+          session = null;
+          return;
+        }
+      }
+
       sessionPool.putBack(session);
       session = null;
     }

diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4
@@ -73,6 +73,8 @@ keyWords
     | CQ
     | CQS
     | CREATE
+    | CURRENT_SQL_DIALECT
+    | CURRENT_USER
     | DATA
     | DATA_REPLICATION_FACTOR
     | DATA_REGION_GROUP_NUM
@@ -214,6 +216,7 @@ keyWords
     | SOFFSET
     | SOURCE
     | SPACE
+    | SQL_DIALECT
     | STORAGE
     | START
     | STARTTIME
@@ -224,6 +227,7 @@ keyWords
     | SUBSCRIPTIONS
     | SUBSTRING
     | SYSTEM
+    | TABLE
     | TAGS
     | TAIL
     | TASK
@@ -242,6 +246,7 @@ keyWords
     | TOPIC
     | TOPICS
     | TRACING
+    | TREE
     | TRIGGER
     | TRIGGERS
     | TRUE

diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4
@@ -88,7 +88,7 @@ utilityStatement
     | setSystemStatus | showVersion | showFlushInfo | showLockInfo | showQueryResource
     | showQueries | showCurrentTimestamp | killQuery | grantWatermarkEmbedding
     | revokeWatermarkEmbedding | loadConfiguration | loadTimeseries | loadFile
-    | removeFile | unloadFile
+    | removeFile | unloadFile | setSqlDialect | showCurrentSqlDialect | showCurrentUser
     ;
 
 /**
@@ -1183,6 +1183,18 @@ unloadFile
     : UNLOAD srcFileName=STRING_LITERAL dstFileDir=STRING_LITERAL
     ;
 
+setSqlDialect
+    : SET SQL_DIALECT OPERATOR_SEQ (TABLE | TREE)
+    ;
+
+showCurrentSqlDialect
+    : SHOW CURRENT_SQL_DIALECT
+    ;
+
+showCurrentUser
+    : SHOW CURRENT_USER
+    ;
+
 // attribute clauses
 syncAttributeClauses
     : attributePair (COMMA? attributePair)*