Skip to content

Commit 148250e

Browse files
CaideyipiCaideyipi
committed
count database
1 parent 402e399 commit 148250e

11 files changed

Lines changed: 251 additions & 1 deletion

File tree

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -825,6 +825,29 @@ public void testMixedDatabase() throws SQLException {
825825
}
826826
}
827827

828+
@Test
829+
public void testCountDatabases() throws SQLException {
830+
try (final Connection connection =
831+
EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
832+
final Statement statement = connection.createStatement()) {
833+
statement.execute("create database db1");
834+
statement.execute("create database db2");
835+
836+
TestUtils.assertResultSetEqual(
837+
statement.executeQuery("count databases"),
838+
"count,",
839+
Collections.singleton("3,"));
840+
TestUtils.assertResultSetEqual(
841+
statement.executeQuery("count database"),
842+
"count,",
843+
Collections.singleton("3,"));
844+
TestUtils.assertResultSetEqual(
845+
statement.executeQuery("select count(*) from information_schema.databases"),
846+
"_col0,",
847+
Collections.singleton("3,"));
848+
}
849+
}
850+
828851
@Test
829852
public void testDBAuth() throws SQLException {
830853
try (final Connection adminCon = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
@@ -838,6 +861,14 @@ public void testDBAuth() throws SQLException {
838861
try (final Connection userCon =
839862
EnvFactory.getEnv().getConnection("test", "password123456", BaseEnv.TABLE_SQL_DIALECT);
840863
final Statement userStmt = userCon.createStatement()) {
864+
TestUtils.assertResultSetEqual(
865+
userStmt.executeQuery("count databases"),
866+
"count,",
867+
Collections.singleton("1,"));
868+
TestUtils.assertResultSetEqual(
869+
userStmt.executeQuery("count database"),
870+
"count,",
871+
Collections.singleton("1,"));
841872
TestUtils.assertResultSetEqual(
842873
userStmt.executeQuery("show databases"),
843874
"Database,TTL(ms),SchemaReplicationFactor,DataReplicationFactor,TimePartitionInterval,",
@@ -869,6 +900,10 @@ public void testDBAuth() throws SQLException {
869900
try (final Connection userCon =
870901
EnvFactory.getEnv().getConnection("test", "password123456", BaseEnv.TABLE_SQL_DIALECT);
871902
final Statement userStmt = userCon.createStatement()) {
903+
TestUtils.assertResultSetEqual(
904+
userStmt.executeQuery("count databases"),
905+
"count,",
906+
Collections.singleton("2,"));
872907
try (final ResultSet resultSet = userStmt.executeQuery("SHOW DATABASES")) {
873908
final ResultSetMetaData metaData = resultSet.getMetaData();
874909
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: 20 additions & 0 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;
@@ -4140,6 +4142,24 @@ public SettableFuture<ConfigTaskResult> showDatabases(
41404142
return future;
41414143
}
41424144

4145+
@Override
4146+
public SettableFuture<ConfigTaskResult> countDatabases(
4147+
final CountDB countDB, final Predicate<String> canSeenDB) {
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, canSeenDB);
4157+
} catch (final IOException | ClientManagerException | TException e) {
4158+
future.setException(e);
4159+
}
4160+
return future;
4161+
}
4162+
41434163
@Override
41444164
public SettableFuture<ConfigTaskResult> showCluster(final ShowCluster showCluster) {
41454165
// As the implementation is identical, we'll simply translate to the

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

Lines changed: 4 additions & 0 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;
@@ -333,6 +334,9 @@ SettableFuture<ConfigTaskResult> showThrottleQuota(
333334
SettableFuture<ConfigTaskResult> showDatabases(
334335
final ShowDB showDB, final Predicate<String> canSeenDB);
335336

337+
SettableFuture<ConfigTaskResult> countDatabases(
338+
final CountDB countDB, final Predicate<String> canSeenDB);
339+
336340
SettableFuture<ConfigTaskResult> showCluster(ShowCluster showCluster);
337341

338342
SettableFuture<ConfigTaskResult> useDatabase(final Use useDB, final IClientSession clientSession);

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
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> canSeenDB;
46+
47+
public CountDBTask(final CountDB node, final Predicate<String> canSeenDB) {
48+
this.node = node;
49+
this.canSeenDB = canSeenDB;
50+
}
51+
52+
@Override
53+
public ListenableFuture<ConfigTaskResult> execute(final IConfigTaskExecutor configTaskExecutor)
54+
throws InterruptedException {
55+
return configTaskExecutor.countDatabases(node, canSeenDB);
56+
}
57+
58+
public static void buildTSBlock(
59+
final Map<String, ?> databaseInfoMap,
60+
final SettableFuture<ConfigTaskResult> future,
61+
final Predicate<String> canSeenDB) {
62+
final long databaseCount =
63+
databaseInfoMap.keySet().stream().filter(canSeenDB::test).count()
64+
+ (canSeenDB.test(INFORMATION_DATABASE) ? 1 : 0);
65+
66+
final TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32));
67+
builder.getTimeColumnBuilder().writeLong(0L);
68+
builder.getColumnBuilder(0).writeInt((int) databaseCount);
69+
builder.declarePosition();
70+
71+
final DatasetHeader datasetHeader =
72+
new DatasetHeader(
73+
Collections.singletonList(
74+
new ColumnHeader(IoTDBConstant.COLUMN_COUNT, TSDataType.INT32)),
75+
true);
76+
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS, builder.build(), datasetHeader));
77+
}
78+
}

‎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
}

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
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.relational.sql.ast;
21+
22+
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.AstMemoryEstimationHelper;
23+
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.IAstVisitor;
24+
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Node;
25+
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.NodeLocation;
26+
import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement;
27+
28+
import com.google.common.collect.ImmutableList;
29+
import org.apache.tsfile.utils.RamUsageEstimator;
30+
31+
import java.util.List;
32+
33+
import static java.util.Objects.requireNonNull;
34+
35+
public class CountDB extends Statement {
36+
37+
private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(CountDB.class);
38+
39+
public CountDB(final NodeLocation location) {
40+
super(requireNonNull(location, "location is null"));
41+
}
42+
43+
@Override
44+
public <R, C> R accept(final IAstVisitor<R, C> visitor, final C context) {
45+
return ((AstVisitor<R, C>) visitor).visitCountDB(this, context);
46+
}
47+
48+
@Override
49+
public List<Node> getChildren() {
50+
return ImmutableList.of();
51+
}
52+
53+
@Override
54+
public int hashCode() {
55+
return getClass().hashCode();
56+
}
57+
58+
@Override
59+
public boolean equals(final Object obj) {
60+
if (this == obj) {
61+
return true;
62+
}
63+
return (obj != null) && (getClass() == obj.getClass());
64+
}
65+
66+
@Override
67+
public String toString() {
68+
return "COUNT DATABASES";
69+
}
70+
71+
@Override
72+
public long ramBytesUsed() {
73+
long size = INSTANCE_SIZE;
74+
size += AstMemoryEstimationHelper.getEstimatedSizeOfNodeLocation(getLocationInternal());
75+
return size;
76+
}
77+
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@
162162
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ClearCache;
163163
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ColumnDefinition;
164164
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CopyTo;
165+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
165166
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDevice;
166167
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountStatement;
167168
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateDB;
@@ -406,6 +407,12 @@ public Node visitShowDatabasesStatement(
406407
return new ShowDB(getLocation(ctx), Objects.nonNull(ctx.DETAILS()));
407408
}
408409

410+
@Override
411+
public Node visitCountDatabasesStatement(
412+
final RelationalSqlParser.CountDatabasesStatementContext ctx) {
413+
return new CountDB(getLocation(ctx));
414+
}
415+
409416
@Override
410417
public Node visitCreateDbStatement(final RelationalSqlParser.CreateDbStatementContext ctx) {
411418
List<Property> properties = ImmutableList.of();

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AstVisitor;
3333
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ColumnDefinition;
3434
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CopyTo;
35+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CountDB;
3536
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateDB;
3637
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreateFunction;
3738
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.CreatePipe;
@@ -134,7 +135,13 @@ public Void visitExplainAnalyze(ExplainAnalyze node, Integer indent) {
134135

135136
@Override
136137
public Void visitShowDB(ShowDB node, Integer indent) {
137-
builder.append("SHOW DATABASE");
138+
builder.append("SHOW DATABASES");
139+
return null;
140+
}
141+
142+
@Override
143+
public Void visitCountDB(CountDB node, Integer indent) {
144+
builder.append("COUNT DATABASES");
138145
return null;
139146
}
140147

0 commit comments

Comments
 (0)