Skip to content

Commit e76945f

Browse files
alex268alex268
committed
Added support of RETURNING statement
1 parent f12bf4e commit e76945f

File tree

3 files changed

+166
-4
lines changed

3 files changed

+166
-4
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,11 @@ public void setHasReturning(boolean hasReturning) {
3535
}
3636

3737
public boolean hasUpdateCount() {
38-
return (command == QueryCmd.INSERT_UPSERT || command == QueryCmd.UPDATE_REPLACE_DELETE); // && !hasReturinng;
38+
return (command == QueryCmd.INSERT_UPSERT || command == QueryCmd.UPDATE_REPLACE_DELETE) && !hasReturinng;
3939
}
4040

4141
public boolean hasResults() {
42-
return command == QueryCmd.SELECT; // || hasReturinng;
42+
return command == QueryCmd.SELECT || hasReturinng;
4343
}
4444

4545
public boolean isDDL() {

‎jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import org.junit.jupiter.api.Assertions;
1717
import org.junit.jupiter.api.BeforeAll;
1818
import org.junit.jupiter.api.BeforeEach;
19-
import org.junit.jupiter.api.Disabled;
2019
import org.junit.jupiter.api.DisplayName;
2120
import org.junit.jupiter.api.Test;
2221
import org.junit.jupiter.api.extension.RegisterExtension;
@@ -686,7 +685,6 @@ public void testMixedStatements() throws SQLException {
686685
}
687686

688687
@Test
689-
@Disabled
690688
public void testReturingStatements() throws SQLException {
691689
String returningQuery = QUERIES.withTableName(""
692690
+ "INSERT INTO #tableName (key, c_Text) VALUES (1, '123') RETURNING key;\n"
@@ -702,6 +700,7 @@ public void testReturingStatements() throws SQLException {
702700
);
703701

704702
try (Statement statement = jdbc.connection().createStatement()) {
703+
// INSERT with returning
705704
Assertions.assertTrue(statement.execute(returningQuery));
706705
Assertions.assertEquals(-1, statement.getUpdateCount());
707706
try (ResultSet rs = statement.getResultSet()) {
@@ -710,7 +709,73 @@ public void testReturingStatements() throws SQLException {
710709
Assertions.assertFalse(rs.next());
711710
}
712711

712+
// simple INSERT
713713
Assertions.assertFalse(statement.getMoreResults());
714+
Assertions.assertNull(statement.getResultSet());
715+
Assertions.assertEquals(1, statement.getUpdateCount());
716+
717+
// UPDATE with returning
718+
Assertions.assertTrue(statement.getMoreResults());
719+
Assertions.assertEquals(-1, statement.getUpdateCount());
720+
try (ResultSet rs = statement.getResultSet()) {
721+
Assertions.assertTrue(rs.next());
722+
Assertions.assertEquals("100", rs.getString("c_Text"));
723+
Assertions.assertFalse(rs.next());
724+
}
725+
726+
// simple UPDATE
727+
Assertions.assertFalse(statement.getMoreResults());
728+
Assertions.assertNull(statement.getResultSet());
729+
Assertions.assertEquals(1, statement.getUpdateCount());
730+
731+
// UPSERT with returning
732+
Assertions.assertTrue(statement.getMoreResults());
733+
Assertions.assertEquals(-1, statement.getUpdateCount());
734+
try (ResultSet rs = statement.getResultSet()) {
735+
Assertions.assertTrue(rs.next());
736+
Assertions.assertEquals(1, rs.getInt("key"));
737+
Assertions.assertEquals("321", rs.getString("c_Text"));
738+
Assertions.assertFalse(rs.next());
739+
}
740+
741+
// simple UPSERT
742+
Assertions.assertFalse(statement.getMoreResults());
743+
Assertions.assertNull(statement.getResultSet());
744+
Assertions.assertEquals(1, statement.getUpdateCount());
745+
746+
// REPLACE with returning
747+
Assertions.assertTrue(statement.getMoreResults());
748+
Assertions.assertEquals(-1, statement.getUpdateCount());
749+
try (ResultSet rs = statement.getResultSet()) {
750+
Assertions.assertTrue(rs.next());
751+
Assertions.assertEquals(1, rs.getInt("key"));
752+
Assertions.assertEquals("111", rs.getString("c_Text"));
753+
Assertions.assertFalse(rs.next());
754+
}
755+
756+
// simple REPLACE
757+
Assertions.assertFalse(statement.getMoreResults());
758+
Assertions.assertNull(statement.getResultSet());
759+
Assertions.assertEquals(1, statement.getUpdateCount());
760+
761+
// DELETE with returning
762+
Assertions.assertTrue(statement.getMoreResults());
763+
Assertions.assertEquals(-1, statement.getUpdateCount());
764+
try (ResultSet rs = statement.getResultSet()) {
765+
Assertions.assertTrue(rs.next());
766+
Assertions.assertEquals("111", rs.getString("c_Text"));
767+
Assertions.assertFalse(rs.next());
768+
}
769+
770+
// simple DELETE
771+
Assertions.assertFalse(statement.getMoreResults());
772+
Assertions.assertNull(statement.getResultSet());
773+
Assertions.assertEquals(1, statement.getUpdateCount());
774+
775+
// no more results
776+
Assertions.assertFalse(statement.getMoreResults());
777+
Assertions.assertNull(statement.getResultSet());
778+
Assertions.assertEquals(-1, statement.getUpdateCount());
714779
}
715780
}
716781

‎jdbc/src/test/java/tech/ydb/jdbc/impl/YdbQueryConnectionImplTest.java

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.junit.jupiter.api.Assertions;
1717
import org.junit.jupiter.api.BeforeAll;
1818
import org.junit.jupiter.api.BeforeEach;
19+
import org.junit.jupiter.api.Disabled;
1920
import org.junit.jupiter.api.DisplayName;
2021
import org.junit.jupiter.api.Test;
2122
import org.junit.jupiter.api.extension.RegisterExtension;
@@ -685,6 +686,102 @@ public void testMixedStatements() throws SQLException {
685686
}
686687
}
687688

689+
@Test
690+
@Disabled // https://github.com/ydb-platform/ydb/issues/6699
691+
public void testReturingStatements() throws SQLException {
692+
String returningQuery = QUERIES.withTableName(""
693+
+ "INSERT INTO #tableName (key, c_Text) VALUES (1, '123') RETURNING key;\n"
694+
+ "INSERT INTO #tableName (key, c_Text) VALUES (2, '234');\n"
695+
+ "UPDATE #tableName SET c_Text = '100' WHERE key = 1 RETURNING c_Text;\n"
696+
+ "UPDATE #tableName SET c_Text = '200' WHERE key = 2;\n"
697+
+ "UPSERT INTO #tableName (key, c_Text) VALUES (1, '321') RETURNING key, c_Text;\n"
698+
+ "UPSERT INTO #tableName (key, c_Text) VALUES (2, '222');\n"
699+
+ "REPLACE INTO #tableName (key, c_Text) VALUES (1, '111') RETURNING c_Text, key;\n"
700+
+ "REPLACE INTO #tableName (key, c_Text) VALUES (2, '333');\n"
701+
+ "DELETE FROM #tableName WHERE key = 1 RETURNING c_Text;\n"
702+
+ "DELETE FROM #tableName WHERE key = 2;\n"
703+
);
704+
705+
try (Statement statement = jdbc.connection().createStatement()) {
706+
// INSERT with returning
707+
Assertions.assertTrue(statement.execute(returningQuery));
708+
Assertions.assertEquals(-1, statement.getUpdateCount());
709+
try (ResultSet rs = statement.getResultSet()) {
710+
Assertions.assertTrue(rs.next());
711+
Assertions.assertEquals(1, rs.getInt("key"));
712+
Assertions.assertFalse(rs.next());
713+
}
714+
715+
// simple INSERT
716+
Assertions.assertFalse(statement.getMoreResults());
717+
Assertions.assertNull(statement.getResultSet());
718+
Assertions.assertEquals(1, statement.getUpdateCount());
719+
720+
// UPDATE with returning
721+
Assertions.assertTrue(statement.getMoreResults());
722+
Assertions.assertEquals(-1, statement.getUpdateCount());
723+
try (ResultSet rs = statement.getResultSet()) {
724+
Assertions.assertTrue(rs.next());
725+
Assertions.assertEquals("100", rs.getString("c_Text"));
726+
Assertions.assertFalse(rs.next());
727+
}
728+
729+
// simple UPDATE
730+
Assertions.assertFalse(statement.getMoreResults());
731+
Assertions.assertNull(statement.getResultSet());
732+
Assertions.assertEquals(1, statement.getUpdateCount());
733+
734+
// UPSERT with returning
735+
Assertions.assertTrue(statement.getMoreResults());
736+
Assertions.assertEquals(-1, statement.getUpdateCount());
737+
try (ResultSet rs = statement.getResultSet()) {
738+
Assertions.assertTrue(rs.next());
739+
Assertions.assertEquals(1, rs.getInt("key"));
740+
Assertions.assertEquals("321", rs.getString("c_Text"));
741+
Assertions.assertFalse(rs.next());
742+
}
743+
744+
// simple UPSERT
745+
Assertions.assertFalse(statement.getMoreResults());
746+
Assertions.assertNull(statement.getResultSet());
747+
Assertions.assertEquals(1, statement.getUpdateCount());
748+
749+
// REPLACE with returning
750+
Assertions.assertTrue(statement.getMoreResults());
751+
Assertions.assertEquals(-1, statement.getUpdateCount());
752+
try (ResultSet rs = statement.getResultSet()) {
753+
Assertions.assertTrue(rs.next());
754+
Assertions.assertEquals(1, rs.getInt("key"));
755+
Assertions.assertEquals("111", rs.getString("c_Text"));
756+
Assertions.assertFalse(rs.next());
757+
}
758+
759+
// simple REPLACE
760+
Assertions.assertFalse(statement.getMoreResults());
761+
Assertions.assertNull(statement.getResultSet());
762+
Assertions.assertEquals(1, statement.getUpdateCount());
763+
764+
// DELETE with returning
765+
Assertions.assertTrue(statement.getMoreResults());
766+
Assertions.assertEquals(-1, statement.getUpdateCount());
767+
try (ResultSet rs = statement.getResultSet()) {
768+
Assertions.assertTrue(rs.next());
769+
Assertions.assertEquals("111", rs.getString("c_Text"));
770+
Assertions.assertFalse(rs.next());
771+
}
772+
773+
// simple DELETE
774+
Assertions.assertFalse(statement.getMoreResults());
775+
Assertions.assertNull(statement.getResultSet());
776+
Assertions.assertEquals(1, statement.getUpdateCount());
777+
778+
// no more results
779+
Assertions.assertFalse(statement.getMoreResults());
780+
Assertions.assertNull(statement.getResultSet());
781+
Assertions.assertEquals(-1, statement.getUpdateCount());
782+
}
783+
}
784+
688785
@Test
689786
public void testWarningsInQuery() throws SQLException {
690787
String createTempTable = QUERIES.withTableName(

0 commit comments

Comments
 (0)