Skip to content

Commit a7344aa

Browse files
alex268alex268
authored
Added flag to forced usage of BulkUpsert and ScanQuery (#66)
2 parents e38aa47 + f5182c6 commit a7344aa

File tree

8 files changed

+742
-0
lines changed

8 files changed

+742
-0
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/query/YdbQuery.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,10 +69,21 @@ public static YdbQuery parseQuery(String query, YdbQueryProperties opts) throws
6969
QueryType type = opts.getForcedQueryType();
7070
if (type == null) {
7171
type = parser.detectQueryType();
72+
73+
if (opts.isForcedScanAndBulks() && type != QueryType.SCHEME_QUERY && type != QueryType.EXPLAIN_QUERY) {
74+
if (parser.getYqlBatcher().isValidBatch()) {
75+
parser.getYqlBatcher().setForcedUpsert();
76+
type = QueryType.BULK_QUERY;
77+
} else {
78+
type = QueryType.SCAN_QUERY;
79+
}
80+
}
7281
}
82+
7383
if (parser.getYqlBatcher().isValidBatch()) {
7484
return new YdbQuery(query, preparedYQL, parser.getStatements(), parser.getYqlBatcher(), type);
7585
}
86+
7687
return new YdbQuery(query, preparedYQL, parser.getStatements(), type);
7788
}
7889
}

‎jdbc/src/main/java/tech/ydb/jdbc/query/YqlBatcher.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,10 @@ private enum State {
3434
private final List<String> columns = new ArrayList<>();
3535
private final List<String> values = new ArrayList<>();
3636

37+
public void setForcedUpsert() {
38+
cmd = Cmd.UPSERT;
39+
}
40+
3741
public boolean isInsert() {
3842
return cmd == Cmd.INSERT;
3943
}

‎jdbc/src/main/java/tech/ydb/jdbc/settings/YdbConfig.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,13 +177,15 @@ public DriverPropertyInfo[] toPropertyInfo() throws SQLException {
177177
YdbOperationProperties.TRANSACTION_LEVEL.toInfo(properties),
178178
YdbOperationProperties.SCHEME_QUERY_TX_MODE.toInfo(properties),
179179
YdbOperationProperties.SCAN_QUERY_TX_MODE.toInfo(properties),
180+
YdbOperationProperties.BULK_QUERY_TX_MODE.toInfo(properties),
180181

181182
YdbQueryProperties.DISABLE_PREPARE_DATAQUERY.toInfo(properties),
182183
YdbQueryProperties.DISABLE_AUTO_PREPARED_BATCHES.toInfo(properties),
183184
YdbQueryProperties.DISABLE_DETECT_SQL_OPERATIONS.toInfo(properties),
184185
YdbQueryProperties.DISABLE_JDBC_PARAMETERS.toInfo(properties),
185186
YdbQueryProperties.DISABLE_JDBC_PARAMETERS_DECLARE.toInfo(properties),
186187
YdbQueryProperties.FORCE_QUERY_MODE.toInfo(properties),
188+
YdbQueryProperties.FORCE_SCAN_BULKS.toInfo(properties),
187189
};
188190
}
189191

‎jdbc/src/main/java/tech/ydb/jdbc/settings/YdbQueryProperties.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@ public class YdbQueryProperties {
3030
static final YdbProperty<QueryType> FORCE_QUERY_MODE = YdbProperty.enums("forceQueryMode", QueryType.class,
3131
"Force usage one of query modes (DATA_QUERY, SCAN_QUERY, SCHEME_QUERY or EXPLAIN_QUERYn) for all statements"
3232
);
33+
static final YdbProperty<Boolean> FORCE_SCAN_BULKS = YdbProperty.bool("forceScanAndBulk",
34+
"Force usage of bulk upserts instead of upserts/inserts and scan query for selects",
35+
false
36+
);
3337

3438
private final boolean isDetectQueryType;
3539
private final boolean isDetectJdbcParameters;
@@ -39,6 +43,7 @@ public class YdbQueryProperties {
3943
private final boolean isDetectBatchQueries;
4044

4145
private final QueryType forcedType;
46+
private final boolean isForcedScanAndBulks;
4247

4348
public YdbQueryProperties(YdbConfig config) throws SQLException {
4449
Properties props = config.getProperties();
@@ -59,6 +64,7 @@ public YdbQueryProperties(YdbConfig config) throws SQLException {
5964
&& !disableJdbcParametersDeclare;
6065

6166
this.forcedType = FORCE_QUERY_MODE.readValue(props).getValue();
67+
this.isForcedScanAndBulks = FORCE_SCAN_BULKS.readValue(props).getValue();
6268
}
6369

6470
public boolean isDetectQueryType() {
@@ -84,4 +90,8 @@ public boolean isDetectBatchQueries() {
8490
public QueryType getForcedQueryType() {
8591
return forcedType;
8692
}
93+
94+
public boolean isForcedScanAndBulks() {
95+
return isForcedScanAndBulks;
96+
}
8797
}

‎jdbc/src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package tech.ydb.jdbc;
22

33
import java.sql.Connection;
4+
import java.sql.Driver;
45
import java.sql.DriverManager;
56
import java.sql.SQLException;
67
import java.util.ArrayDeque;
@@ -29,6 +30,25 @@ public class YdbDriverIntegrationTest {
2930

3031
private static final JdbcUrlHelper jdbcURL = new JdbcUrlHelper(ydb);
3132

33+
@Test
34+
public void registerTest() throws SQLException {
35+
try {
36+
YdbDriver.deregister();
37+
} catch (IllegalStateException ex) {
38+
// skip error if driver was deregistered already
39+
}
40+
41+
IllegalStateException ex1 = Assertions.assertThrows(IllegalStateException.class, () -> YdbDriver.deregister());
42+
Assertions.assertEquals(YdbConst.DRIVER_IS_NOT_REGISTERED, ex1.getMessage());
43+
44+
Assertions.assertFalse(YdbDriver.isRegistered());
45+
YdbDriver.register();
46+
Assertions.assertTrue(YdbDriver.isRegistered());
47+
48+
IllegalStateException ex2 = Assertions.assertThrows(IllegalStateException.class, () -> YdbDriver.register());
49+
Assertions.assertEquals(YdbConst.DRIVER_IS_ALREADY_REGISTERED, ex2.getMessage());
50+
}
51+
3252
@Test
3353
public void connect() throws SQLException {
3454
try (Connection connection = DriverManager.getConnection(jdbcURL.build())) {
@@ -43,8 +63,15 @@ public void connect() throws SQLException {
4363

4464
@Test
4565
public void testContextCache() throws SQLException {
66+
Driver jdbcDriver = DriverManager.getDriver(jdbcURL.build());
67+
Assertions.assertTrue(jdbcDriver instanceof YdbDriver);
68+
69+
YdbDriver driver = (YdbDriver)jdbcDriver;
70+
Assertions.assertEquals(0, driver.getConnectionCount());
71+
4672
Connection firstConnection = DriverManager.getConnection(jdbcURL.build());
4773
Assertions.assertTrue(firstConnection.isValid(5000));
74+
Assertions.assertEquals(1, driver.getConnectionCount());
4875

4976
YdbConnection first = firstConnection.unwrap(YdbConnection.class);
5077
Assertions.assertNotNull(first.getCtx());
@@ -53,6 +80,7 @@ public void testContextCache() throws SQLException {
5380

5481
try (Connection conn2 = DriverManager.getConnection(jdbcURL.build())) {
5582
Assertions.assertTrue(conn2.isValid(5000));
83+
Assertions.assertEquals(1, driver.getConnectionCount());
5684

5785
YdbConnection unwrapped = conn2.unwrap(YdbConnection.class);
5886
Assertions.assertNotNull(unwrapped.getCtx());
@@ -62,6 +90,7 @@ public void testContextCache() throws SQLException {
6290
Properties props = new Properties();
6391
try (Connection conn3 = DriverManager.getConnection(jdbcURL.build(), props)) {
6492
Assertions.assertTrue(conn3.isValid(5000));
93+
Assertions.assertEquals(1, driver.getConnectionCount());
6594

6695
YdbConnection unwrapped = conn3.unwrap(YdbConnection.class);
6796
Assertions.assertNotNull(unwrapped.getCtx());
@@ -71,6 +100,7 @@ public void testContextCache() throws SQLException {
71100
props.setProperty("TEST_KEY", "TEST_VALUE");
72101
try (Connection conn4 = DriverManager.getConnection(jdbcURL.build(), props)) {
73102
Assertions.assertTrue(conn4.isValid(5000));
103+
Assertions.assertEquals(2, driver.getConnectionCount());
74104

75105
YdbConnection unwrapped = conn4.unwrap(YdbConnection.class);
76106
Assertions.assertNotNull(unwrapped.getCtx());
@@ -79,21 +109,64 @@ public void testContextCache() throws SQLException {
79109

80110
try (Connection conn5 = DriverManager.getConnection(jdbcURL.withArg("test", "false").build())) {
81111
Assertions.assertTrue(conn5.isValid(5000));
112+
Assertions.assertEquals(2, driver.getConnectionCount());
82113

83114
YdbConnection unwrapped = conn5.unwrap(YdbConnection.class);
84115
Assertions.assertNotNull(unwrapped.getCtx());
85116
Assertions.assertNotSame(ctx, unwrapped.getCtx());
86117
}
87118

88119
firstConnection.close();
120+
Assertions.assertEquals(0, driver.getConnectionCount());
89121

90122
try (Connection conn6 = DriverManager.getConnection(jdbcURL.build())) {
91123
Assertions.assertTrue(conn6.isValid(5000));
124+
Assertions.assertEquals(1, driver.getConnectionCount());
92125

93126
YdbConnection unwrapped = conn6.unwrap(YdbConnection.class);
94127
Assertions.assertNotNull(unwrapped.getCtx());
95128
Assertions.assertNotSame(ctx, unwrapped.getCtx());
96129
}
130+
131+
Assertions.assertEquals(0, driver.getConnectionCount());
132+
}
133+
134+
@Test
135+
public void testContextCacheDisable() throws SQLException {
136+
Driver jdbcDriver = DriverManager.getDriver(jdbcURL.build());
137+
Assertions.assertTrue(jdbcDriver instanceof YdbDriver);
138+
YdbDriver driver = (YdbDriver)jdbcDriver;
139+
Assertions.assertEquals(0, driver.getConnectionCount());
140+
141+
Properties props = new Properties();
142+
props.put("cacheConnectionsInDriver", "false");
143+
144+
Connection conn1 = DriverManager.getConnection(jdbcURL.build(), props);
145+
Assertions.assertEquals(0, driver.getConnectionCount());
146+
147+
YdbConnection first = conn1.unwrap(YdbConnection.class);
148+
Assertions.assertNotNull(first.getCtx());
149+
YdbContext ctx = first.getCtx();
150+
151+
try (Connection conn2 = DriverManager.getConnection(jdbcURL.build())) {
152+
Assertions.assertEquals(1, driver.getConnectionCount());
153+
154+
YdbConnection second = conn2.unwrap(YdbConnection.class);
155+
Assertions.assertNotNull(second.getCtx());
156+
Assertions.assertNotSame(ctx, second.getCtx());
157+
}
158+
159+
Assertions.assertEquals(0, driver.getConnectionCount());
160+
161+
try (Connection conn2 = DriverManager.getConnection(jdbcURL.build(), props)) {
162+
Assertions.assertEquals(0, driver.getConnectionCount());
163+
164+
YdbConnection second = conn2.unwrap(YdbConnection.class);
165+
Assertions.assertNotNull(second.getCtx());
166+
Assertions.assertNotSame(ctx, second.getCtx());
167+
}
168+
169+
conn1.close();
97170
}
98171

99172
@Test

0 commit comments

Comments
 (0)