Skip to content

Commit 4a69937

Browse files
shizy818shizy818
committed
feat: support automatic schema creation in Session
1 parent d2b97bc commit 4a69937

12 files changed

Lines changed: 129 additions & 5 deletions

File tree

‎integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBSessionRelationalIT.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,45 @@ public void partialInsertSQLTest() throws IoTDBConnectionException, StatementExe
456456
}
457457
}
458458

459+
@Test
460+
public void disableAutoCreateSchemaForTableSessionTest()
461+
throws IoTDBConnectionException, StatementExecutionException {
462+
final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0);
463+
try (ITableSession session =
464+
new TableSessionBuilder()
465+
.nodeUrls(Collections.singletonList(dataNode.getIpAndPortString()))
466+
.database("db1")
467+
.enableAutoCreateSchema(false)
468+
.build()) {
469+
final List<IMeasurementSchema> schemaList = new ArrayList<>();
470+
schemaList.add(new MeasurementSchema("tag1", TSDataType.STRING));
471+
schemaList.add(new MeasurementSchema("s1", TSDataType.INT64));
472+
schemaList.add(new MeasurementSchema("s2", TSDataType.INT64));
473+
final Tablet tablet =
474+
new Tablet(
475+
"session_auto_create",
476+
IMeasurementSchema.getMeasurementNameList(schemaList),
477+
IMeasurementSchema.getDataTypeList(schemaList),
478+
Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD),
479+
1);
480+
tablet.addTimestamp(0, 1);
481+
tablet.addValue("tag1", 0, "d1");
482+
tablet.addValue("s1", 0, 1L);
483+
tablet.addValue("s2", 0, 2L);
484+
485+
assertThrows(StatementExecutionException.class, () -> session.insert(tablet));
486+
487+
session.executeNonQueryStatement(
488+
"CREATE TABLE session_auto_create (tag1 string tag, s1 int64 field)");
489+
assertThrows(StatementExecutionException.class, () -> session.insert(tablet));
490+
491+
try (SessionDataSet dataSet =
492+
session.executeQueryStatement("SELECT * FROM session_auto_create")) {
493+
assertEquals(Arrays.asList("time", "tag1", "s1"), dataSet.getColumnNames());
494+
}
495+
}
496+
}
497+
459498
@Test
460499
public void partialInsertRelationalTabletTest()
461500
throws IoTDBConnectionException, StatementExecutionException {

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ public abstract class AbstractSessionBuilder {
8080

8181
public String database;
8282

83+
public boolean enableAutoCreateSchema = true;
84+
8385
public Boolean isIoTDBRpcCompressionEnabled = true;
8486

8587
public Boolean isThriftRpcCompactionEnabled = false;

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,8 @@ public class Session implements ISession {
204204
// may be null
205205
protected volatile String database;
206206

207+
protected volatile boolean enableAutoCreateSchema = true;
208+
207209
private static final String VALUES_SIZE_SHOULD_BE_EQUAL =
208210
"times, measurementsList and valuesList's size should be equal";
209211

@@ -484,6 +486,7 @@ public Session(AbstractSessionBuilder builder) {
484486
this.sqlDialect = builder.sqlDialect;
485487
this.queryTimeoutInMs = builder.timeOut;
486488
this.database = builder.database;
489+
this.enableAutoCreateSchema = builder.enableAutoCreateSchema;
487490
}
488491

489492
@Override
@@ -4523,6 +4526,11 @@ public Builder database(String database) {
45234526
return this;
45244527
}
45254528

4529+
public Builder enableAutoCreateSchema(boolean enableAutoCreateSchema) {
4530+
this.enableAutoCreateSchema = enableAutoCreateSchema;
4531+
return this;
4532+
}
4533+
45264534
public Builder useSSL(boolean useSSL) {
45274535
this.useSSL = useSSL;
45284536
return this;

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ public class SessionConnection {
9999

100100
private static final Logger logger = LoggerFactory.getLogger(SessionConnection.class);
101101
private static final String USE_ENCRYPTED_PASSWORD_KEY = "use_encrypted_password";
102+
private static final String ENABLE_AUTO_CREATE_SCHEMA_KEY = "enable_auto_create_schema";
102103
public static final String MSG_RECONNECTION_FAIL =
103104
"Fail to reconnect to server. Please check server status.";
104105
protected Session session;
@@ -244,6 +245,9 @@ private void init(
244245
if (session.useEncryptedPassword) {
245246
openReq.putToConfiguration(USE_ENCRYPTED_PASSWORD_KEY, Boolean.TRUE.toString());
246247
}
248+
if (!session.enableAutoCreateSchema) {
249+
openReq.putToConfiguration(ENABLE_AUTO_CREATE_SCHEMA_KEY, Boolean.FALSE.toString());
250+
}
247251
if (database != null) {
248252
openReq.putToConfiguration("db", database);
249253
}

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,18 @@ public TableSessionBuilder enableAutoFetch(boolean enableAutoFetch) {
179179
return this;
180180
}
181181

182+
/**
183+
* Enables or disables schema auto-creation for this session's write requests.
184+
*
185+
* @param enableAutoCreateSchema whether to allow auto creating tables and columns.
186+
* @return the current {@link TableSessionBuilder} instance.
187+
* @defaultValue true
188+
*/
189+
public TableSessionBuilder enableAutoCreateSchema(boolean enableAutoCreateSchema) {
190+
this.enableAutoCreateSchema = enableAutoCreateSchema;
191+
return this;
192+
}
193+
182194
/**
183195
* Sets the maximum number of retries for connection attempts.
184196
*

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,8 @@ public class SessionPool implements ISessionPool {
188188
// may be null
189189
protected String database;
190190

191+
protected boolean enableAutoCreateSchema = true;
192+
191193
private static final String INSERT_RECORD_FAIL = SessionMessages.INSERT_RECORD_FAILED;
192194

193195
private static final String INSERT_RECORD_ERROR_MSG =
@@ -553,6 +555,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) {
553555
this.sqlDialect = builder.sqlDialect;
554556
this.database = builder.database;
555557
this.queryTimeoutInMs = builder.timeOut;
558+
this.enableAutoCreateSchema = builder.enableAutoCreateSchema;
556559

557560
if (enableAutoFetch) {
558561
initThreadPool();
@@ -612,6 +615,7 @@ private Session constructNewSession() {
612615
.retryIntervalInMs(retryIntervalInMs)
613616
.sqlDialect(sqlDialect)
614617
.database(database)
618+
.enableAutoCreateSchema(enableAutoCreateSchema)
615619
.timeOut(queryTimeoutInMs)
616620
.enableIoTDBRpcCompression(enableIoTDBRpcCompression)
617621
.enableThriftRpcCompression(enableThriftCompression)
@@ -640,6 +644,7 @@ private Session constructNewSession() {
640644
.retryIntervalInMs(retryIntervalInMs)
641645
.sqlDialect(sqlDialect)
642646
.database(database)
647+
.enableAutoCreateSchema(enableAutoCreateSchema)
643648
.timeOut(queryTimeoutInMs)
644649
.enableIoTDBRpcCompression(enableIoTDBRpcCompression)
645650
.enableThriftRpcCompression(enableThriftCompression)

‎iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,18 @@ public TableSessionPoolBuilder sslProtocol(String sslProtocol) {
317317
return this;
318318
}
319319

320+
/**
321+
* Enables or disables schema auto-creation for this pool's write requests.
322+
*
323+
* @param enableAutoCreateSchema whether to allow auto creating tables and columns.
324+
* @return the current {@link TableSessionPoolBuilder} instance.
325+
* @defaultValue true
326+
*/
327+
public TableSessionPoolBuilder enableAutoCreateSchema(boolean enableAutoCreateSchema) {
328+
this.enableAutoCreateSchema = enableAutoCreateSchema;
329+
return this;
330+
}
331+
320332
/**
321333
* Builds and returns a configured {@link ITableSessionPool} instance.
322334
*

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ public abstract class IClientSession {
5959

6060
@Nullable private String databaseName;
6161

62+
private boolean enableAutoCreateSchema = true;
63+
6264
private long lastActiveTime = CommonDateTimeUtils.currentTime();
6365

6466
public abstract String getClientAddress();
@@ -201,6 +203,14 @@ public void setDatabaseName(@Nullable String databaseName) {
201203
}
202204
}
203205

206+
public boolean isEnableAutoCreateSchema() {
207+
return enableAutoCreateSchema;
208+
}
209+
210+
public void setEnableAutoCreateSchema(boolean enableAutoCreateSchema) {
211+
this.enableAutoCreateSchema = enableAutoCreateSchema;
212+
}
213+
204214
/**
205215
* Add a prepared statement to this session.
206216
*

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -510,7 +510,8 @@ public SessionInfo getSessionInfo(IClientSession session) {
510510
session.getZoneId(),
511511
session.getClientVersion(),
512512
session.getDatabaseName(),
513-
session.getSqlDialect());
513+
session.getSqlDialect(),
514+
session.isEnableAutoCreateSchema());
514515
}
515516

516517
// Sometimes we need to switch from table model to tree model,

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,7 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
270270

271271
public static final String ERROR_CODE = DataNodeMiscMessages.ERROR_CODE;
272272
private static final String USE_ENCRYPTED_PASSWORD_KEY = "use_encrypted_password";
273+
private static final String ENABLE_AUTO_CREATE_SCHEMA_KEY = "enable_auto_create_schema";
273274

274275
private static final TSProtocolVersion CURRENT_RPC_VERSION =
275276
TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3;
@@ -1633,8 +1634,11 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException {
16331634
parseUseEncryptedPassword(req));
16341635
TSStatus tsStatus = RpcUtils.getStatus(openSessionResp.getCode(), openSessionResp.getMessage());
16351636

1636-
if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() && database.isPresent()) {
1637-
clientSession.setDatabaseName(database.get().toLowerCase(Locale.ENGLISH));
1637+
if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
1638+
if (database.isPresent()) {
1639+
clientSession.setDatabaseName(database.get().toLowerCase(Locale.ENGLISH));
1640+
}
1641+
clientSession.setEnableAutoCreateSchema(parseEnableAutoCreateSchema(req));
16381642
}
16391643
TSOpenSessionResp resp = new TSOpenSessionResp(tsStatus, CURRENT_RPC_VERSION);
16401644
Map<String, String> configuration = new HashMap<>();
@@ -1678,6 +1682,13 @@ private boolean parseUseEncryptedPassword(TSOpenSessionReq req) {
16781682
&& Boolean.parseBoolean(configuration.get(USE_ENCRYPTED_PASSWORD_KEY));
16791683
}
16801684

1685+
private boolean parseEnableAutoCreateSchema(TSOpenSessionReq req) {
1686+
Map<String, String> configuration = req.configuration;
1687+
return configuration == null
1688+
|| !configuration.containsKey(ENABLE_AUTO_CREATE_SCHEMA_KEY)
1689+
|| Boolean.parseBoolean(configuration.get(ENABLE_AUTO_CREATE_SCHEMA_KEY));
1690+
}
1691+
16811692
@Override
16821693
public TSStatus closeSession(TSCloseSessionReq req) {
16831694
return new TSStatus(

0 commit comments

Comments
 (0)