Skip to content

Commit 0c25e53

Browse files
CaideyipiCaideyipi
authored
Implemented the count database in table model (#17705)
1 parent dcb630f commit 0c25e53

14 files changed

Lines changed: 434 additions & 16 deletions

File tree

‎integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -816,6 +816,12 @@ public void testMixedDatabase() throws SQLException {
816816
try (final Connection connection =
817817
EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
818818
final Statement statement = connection.createStatement()) {
819+
TestUtils.assertResultSetEqual(
820+
statement.executeQuery("count databases"), "count,", Collections.singleton("3,"));
821+
TestUtils.assertResultSetEqual(
822+
statement.executeQuery("select count(*) from information_schema.databases"),
823+
"_col0,",
824+
Collections.singleton("3,"));
819825
statement.execute("drop database test");
820826
}
821827

@@ -826,6 +832,23 @@ public void testMixedDatabase() throws SQLException {
826832
}
827833
}
828834

835+
@Test
836+
public void testCountDatabases() throws SQLException {
837+
try (final Connection connection =
838+
EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
839+
final Statement statement = connection.createStatement()) {
840+
statement.execute("create database db1");
841+
statement.execute("create database db2");
842+
843+
TestUtils.assertResultSetEqual(
844+
statement.executeQuery("count databases"), "count,", Collections.singleton("3,"));
845+
TestUtils.assertResultSetEqual(
846+
statement.executeQuery("select count(*) from information_schema.databases"),
847+
"_col0,",
848+
Collections.singleton("3,"));
849+
}
850+
}
851+
829852
@Test
830853
public void testDBAuth() throws SQLException {
831854
try (final Connection adminCon = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
@@ -839,6 +862,8 @@ public void testDBAuth() throws SQLException {
839862
try (final Connection userCon =
840863
EnvFactory.getEnv().getConnection("test", "password123456", BaseEnv.TABLE_SQL_DIALECT);
841864
final Statement userStmt = userCon.createStatement()) {
865+
TestUtils.assertResultSetEqual(
866+
userStmt.executeQuery("count databases"), "count,", Collections.singleton("1,"));
842867
TestUtils.assertResultSetEqual(
843868
userStmt.executeQuery("show databases"),
844869
"Database,TTL(ms),SchemaReplicationFactor,DataReplicationFactor,TimePartitionInterval,",
@@ -870,6 +895,8 @@ public void testDBAuth() throws SQLException {
870895
try (final Connection userCon =
871896
EnvFactory.getEnv().getConnection("test", "password123456", BaseEnv.TABLE_SQL_DIALECT);
872897
final Statement userStmt = userCon.createStatement()) {
898+
TestUtils.assertResultSetEqual(
899+
userStmt.executeQuery("count databases"), "count,", Collections.singleton("2,"));
873900
try (final ResultSet resultSet = userStmt.executeQuery("SHOW DATABASES")) {
874901
final ResultSetMetaData metaData = resultSet.getMetaData();
875902
assertEquals(showDBColumnHeaders.size(), metaData.getColumnCount());

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@
7878
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AlterColumnDataType;
7979
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AlterDB;
8080
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ClearCache;
81+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
8182
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateDB;
8283
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateExternalService;
8384
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateFunction;
@@ -610,6 +611,7 @@ private IQueryExecution createQueryExecutionForTableModel(
610611
queryContext.setTimeOut(timeOut);
611612
queryContext.setStartTime(startTime);
612613
if (statement instanceof DropDB
614+
|| statement instanceof CountDB
613615
|| statement instanceof ShowDB
614616
|| statement instanceof CreateDB
615617
|| statement instanceof AlterDB

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@
9999
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.AlterTableRenameTableTask;
100100
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.AlterTableSetPropertiesTask;
101101
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.ClearCacheTask;
102+
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.CountDBTask;
102103
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.CreateDBTask;
103104
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.CreateTableTask;
104105
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.CreateTableViewTask;
@@ -156,6 +157,7 @@
156157
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AstVisitor;
157158
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ClearCache;
158159
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ColumnDefinition;
160+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
159161
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateDB;
160162
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateExternalService;
161163
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateFunction;
@@ -424,6 +426,15 @@ public IConfigTask visitShowDB(final ShowDB node, final MPPQueryContext context)
424426
canShowDB(accessControl, context.getSession().getUserName(), databaseName, context));
425427
}
426428

429+
@Override
430+
public IConfigTask visitCountDB(final CountDB node, final MPPQueryContext context) {
431+
context.setQueryType(QueryType.READ);
432+
return new CountDBTask(
433+
node,
434+
databaseName ->
435+
canShowDB(accessControl, context.getSession().getUserName(), databaseName, context));
436+
}
437+
427438
public static boolean canShowDB(
428439
final AccessControl accessControl,
429440
final String userName,

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@
229229
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.region.MigrateRegionTask;
230230
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.region.ReconstructRegionTask;
231231
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.region.RemoveRegionTask;
232+
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.CountDBTask;
232233
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.DeleteDeviceTask;
233234
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.DescribeTableDetailsTask;
234235
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.DescribeTableTask;
@@ -255,6 +256,7 @@
255256
import org.apache.iotdb.db.queryengine.plan.expression.Expression;
256257
import org.apache.iotdb.db.queryengine.plan.expression.visitor.TransformToViewExpressionVisitor;
257258
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.view.AlterLogicalViewNode;
259+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
258260
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DeleteDevice;
259261
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DropDB;
260262
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ShowCluster;
@@ -4121,7 +4123,7 @@ public TPipeTransferResp handleTransferConfigPlan(
41214123

41224124
@Override
41234125
public SettableFuture<ConfigTaskResult> showDatabases(
4124-
final ShowDB showDB, final Predicate<String> canSeenDB) {
4126+
final ShowDB showDB, final Predicate<String> canSeeDB) {
41254127
final SettableFuture<ConfigTaskResult> future = SettableFuture.create();
41264128
// Construct request using statement
41274129
final List<String> databasePathPattern = Arrays.asList(ALL_RESULT_NODES);
@@ -4133,7 +4135,25 @@ public SettableFuture<ConfigTaskResult> showDatabases(
41334135
.setIsTableModel(true);
41344136
final TShowDatabaseResp resp = client.showDatabase(req);
41354137
// build TSBlock
4136-
ShowDBTask.buildTSBlock(resp.getDatabaseInfoMap(), future, showDB.isDetails(), canSeenDB);
4138+
ShowDBTask.buildTSBlock(resp.getDatabaseInfoMap(), future, showDB.isDetails(), canSeeDB);
4139+
} catch (final IOException | ClientManagerException | TException e) {
4140+
future.setException(e);
4141+
}
4142+
return future;
4143+
}
4144+
4145+
@Override
4146+
public SettableFuture<ConfigTaskResult> countDatabases(
4147+
final CountDB countDB, final Predicate<String> canSeeDB) {
4148+
final SettableFuture<ConfigTaskResult> future = SettableFuture.create();
4149+
final List<String> databasePathPattern = Arrays.asList(ALL_RESULT_NODES);
4150+
try (final ConfigNodeClient client =
4151+
CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
4152+
final TGetDatabaseReq req =
4153+
new TGetDatabaseReq(databasePathPattern, ALL_MATCH_SCOPE.serialize())
4154+
.setIsTableModel(true);
4155+
final TShowDatabaseResp resp = client.showDatabase(req);
4156+
CountDBTask.buildTSBlock(resp.getDatabaseInfoMap(), future, canSeeDB);
41374157
} catch (final IOException | ClientManagerException | TException e) {
41384158
future.setException(e);
41394159
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.region.ReconstructRegionTask;
4242
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.region.RemoveRegionTask;
4343
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.view.AlterLogicalViewNode;
44+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
4445
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DeleteDevice;
4546
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DropDB;
4647
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ShowCluster;
@@ -331,7 +332,10 @@ SettableFuture<ConfigTaskResult> showThrottleQuota(
331332
// =============================== table syntax =========================================
332333

333334
SettableFuture<ConfigTaskResult> showDatabases(
334-
final ShowDB showDB, final Predicate<String> canSeenDB);
335+
final ShowDB showDB, final Predicate<String> canSeeDB);
336+
337+
SettableFuture<ConfigTaskResult> countDatabases(
338+
final CountDB countDB, final Predicate<String> canSeeDB);
335339

336340
SettableFuture<ConfigTaskResult> showCluster(ShowCluster showCluster);
337341

@@ -355,7 +359,7 @@ SettableFuture<ConfigTaskResult> describeTable(
355359
final Boolean isShowCreateView);
356360

357361
SettableFuture<ConfigTaskResult> showTables(
358-
final String database, final Predicate<String> canSeenDB, final boolean isDetails);
362+
final String database, final Predicate<String> checkCanShowTable, final boolean isDetails);
359363

360364
TFetchTableResp fetchTables(final Map<String, Set<String>> fetchTableMap);
361365

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/relational/CountDBTask.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational;
21+
22+
import org.apache.iotdb.commons.conf.IoTDBConstant;
23+
import org.apache.iotdb.commons.schema.column.ColumnHeader;
24+
import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
25+
import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
26+
import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask;
27+
import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor;
28+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
29+
import org.apache.iotdb.rpc.TSStatusCode;
30+
31+
import com.google.common.util.concurrent.ListenableFuture;
32+
import com.google.common.util.concurrent.SettableFuture;
33+
import org.apache.tsfile.enums.TSDataType;
34+
import org.apache.tsfile.read.common.block.TsBlockBuilder;
35+
36+
import java.util.Collections;
37+
import java.util.Map;
38+
import java.util.function.Predicate;
39+
40+
import static org.apache.iotdb.commons.schema.table.InformationSchema.INFORMATION_DATABASE;
41+
42+
public class CountDBTask implements IConfigTask {
43+
44+
private final CountDB node;
45+
private final Predicate<String> canSeeDB;
46+
47+
public CountDBTask(final CountDB node, final Predicate<String> canSeeDB) {
48+
this.node = node;
49+
this.canSeeDB = canSeeDB;
50+
}
51+
52+
@Override
53+
public ListenableFuture<ConfigTaskResult> execute(final IConfigTaskExecutor configTaskExecutor)
54+
throws InterruptedException {
55+
return configTaskExecutor.countDatabases(node, canSeeDB);
56+
}
57+
58+
public static void buildTSBlock(
59+
final Map<String, ?> databaseInfoMap,
60+
final SettableFuture<ConfigTaskResult> future,
61+
final Predicate<String> canSeeDB) {
62+
// information_schema is synthesized in table model rather than returned from ConfigNode.
63+
final long databaseCount =
64+
databaseInfoMap.keySet().stream()
65+
.filter(databaseName -> !INFORMATION_DATABASE.equals(databaseName))
66+
.filter(canSeeDB::test)
67+
.count()
68+
+ (canSeeDB.test(INFORMATION_DATABASE) ? 1 : 0);
69+
70+
final TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32));
71+
builder.getTimeColumnBuilder().writeLong(0L);
72+
builder.getColumnBuilder(0).writeInt((int) databaseCount);
73+
builder.declarePosition();
74+
75+
final DatasetHeader datasetHeader =
76+
new DatasetHeader(
77+
Collections.singletonList(
78+
new ColumnHeader(IoTDBConstant.COLUMN_COUNT, TSDataType.INT32)),
79+
true);
80+
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS, builder.build(), datasetHeader));
81+
}
82+
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/relational/ShowDBTask.java‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -51,35 +51,35 @@ public class ShowDBTask implements IConfigTask {
5151
private final ShowDB node;
5252

5353
// judge whether the specific database can be seen, dbName should be without `root.` prefix
54-
private final Predicate<String> canSeenDB;
54+
private final Predicate<String> canSeeDB;
5555

56-
public ShowDBTask(final ShowDB node, final Predicate<String> canSeenDB) {
56+
public ShowDBTask(final ShowDB node, final Predicate<String> canSeeDB) {
5757
this.node = node;
58-
this.canSeenDB = canSeenDB;
58+
this.canSeeDB = canSeeDB;
5959
}
6060

6161
@Override
6262
public ListenableFuture<ConfigTaskResult> execute(final IConfigTaskExecutor configTaskExecutor)
6363
throws InterruptedException {
64-
return configTaskExecutor.showDatabases(node, canSeenDB);
64+
return configTaskExecutor.showDatabases(node, canSeeDB);
6565
}
6666

6767
public static void buildTSBlock(
6868
final Map<String, TDatabaseInfo> databaseInfoMap,
6969
final SettableFuture<ConfigTaskResult> future,
7070
final boolean isDetails,
71-
final Predicate<String> canSeenDB) {
71+
final Predicate<String> canSeeDB) {
7272
if (isDetails) {
73-
buildTSBlockForDetails(databaseInfoMap, future, canSeenDB);
73+
buildTSBlockForDetails(databaseInfoMap, future, canSeeDB);
7474
} else {
75-
buildTSBlockForNonDetails(databaseInfoMap, future, canSeenDB);
75+
buildTSBlockForNonDetails(databaseInfoMap, future, canSeeDB);
7676
}
7777
}
7878

7979
private static void buildTSBlockForNonDetails(
8080
final Map<String, TDatabaseInfo> databaseInfoMap,
8181
final SettableFuture<ConfigTaskResult> future,
82-
final Predicate<String> canSeenDB) {
82+
final Predicate<String> canSeeDB) {
8383
final List<TSDataType> outputDataTypes =
8484
ColumnHeaderConstant.showDBColumnHeaders.stream()
8585
.map(ColumnHeader::getColumnType)
@@ -93,7 +93,7 @@ private static void buildTSBlockForNonDetails(
9393
.sorted(Map.Entry.comparingByKey())
9494
.collect(Collectors.toList())) {
9595
final String dbName = entry.getKey();
96-
if (Boolean.FALSE.equals(canSeenDB.test(dbName))) {
96+
if (Boolean.FALSE.equals(canSeeDB.test(dbName))) {
9797
continue;
9898
}
9999
if (dbName.equals(INFORMATION_DATABASE)) {
@@ -127,7 +127,7 @@ private static void buildTSBlockForNonDetails(
127127
private static void buildTSBlockForDetails(
128128
final Map<String, TDatabaseInfo> databaseMap,
129129
final SettableFuture<ConfigTaskResult> future,
130-
final Predicate<String> canSeenDB) {
130+
final Predicate<String> canSeeDB) {
131131
final List<TSDataType> outputDataTypes =
132132
ColumnHeaderConstant.showDBDetailsColumnHeaders.stream()
133133
.map(ColumnHeader::getColumnType)
@@ -141,7 +141,7 @@ private static void buildTSBlockForDetails(
141141
.sorted(Map.Entry.comparingByKey())
142142
.collect(Collectors.toList())) {
143143
final String dbName = entry.getKey();
144-
if (!canSeenDB.test(dbName)) {
144+
if (!canSeeDB.test(dbName)) {
145145
continue;
146146
}
147147
if (dbName.equals(INFORMATION_DATABASE)) {

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/AstVisitor.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@ default R visitShowDB(final ShowDB node, final C context) {
7979
return visitStatement(node, context);
8080
}
8181

82+
default R visitCountDB(final CountDB node, final C context) {
83+
return visitStatement(node, context);
84+
}
85+
8286
default R visitCreateTable(final CreateTable node, final C context) {
8387
return visitStatement(node, context);
8488
}

0 commit comments

Comments
 (0)