Statement batch doesn't produce the correct number of update counts

Bug Report

Versions

• Driver: 0.9.0.M1
• Database: Microsoft SQL Server 2019 (RTM-CU8-GDR) (KB4583459) - 15.0.4083.2 (X64) Nov 2 2020 18:35:09 Copyright © 2019 Microsoft Corporation Express Edition (64-bit) on Linux (Ubuntu 18.04.5 LTS) <X64>
• Java: openjdk version “11.0.9.1” 2020-11-04
• OS: Microsoft Windows [Version 10.0.19042.867]

Current Behavior

This code:

System.out.println((
    Flux.from(connectionFactory.create())
        .flatMap(c -> c
            .createStatement("DECLARE @t TABLE(i INT);\n"
                + "INSERT INTO @t VALUES (1),(2),(3);\n"
                + "SELECT * FROM @t;\n")
            .execute())
        .flatMap(it -> {
            System.out.println("IT: " + it);
            return it.getRowsUpdated();
        })
        .collectList()
        .block()
));

Prints:

IT: io.r2dbc.mssql.MssqlResult@3df5a869
[6]

Expected behavior/code

I would expect the same behaviour as in JDBC. Here’s a loop that collects both update counts and result sets:

try (Statement s = connection.createStatement()) {
    boolean b = s.execute("DECLARE @t TABLE(i INT);\n"
        + "INSERT INTO @t VALUES (1),(2),(3);\n"
        + "SELECT * FROM @t;\n");
    while (true) {
        int c = -1;

        if (b) {
            try (ResultSet rs = s.getResultSet()) {
                while (rs.next())
                    System.out.println("rs: " + rs.getInt(1));
            }
        }
        else
            c = s.getUpdateCount();

        System.out.println(b + ": " + c);

        if (!b && c == -1)
            break;
        else
            b = s.getMoreResults();
    }
}

Printing:

false: 3
rs: 1
rs: 2
rs: 3
true: -1
false: -1

The JDBC version actually checks whether the current result is an update count or a result set (this isn’t currently possible with R2DBC, see https://github.com/r2dbc/r2dbc-spi/issues/27). But even if we embed this knowledge into the reactive version, I’m not getting the right results:

AtomicInteger i = new AtomicInteger(0);
System.out.println((
    Flux.from(connectionFactory.create())
        .flatMap(c -> c
            .createStatement("DECLARE @t TABLE(i INT);\n"
                + "INSERT INTO @t VALUES (1),(2),(3);\n"
                + "SELECT * FROM @t;\n")
            .execute())
        .flatMap(it -> {
            System.out.println("IT: " + it);
            if (i.getAndIncrement() == 0)
                return it.getRowsUpdated();
            else
                return it.map((r, m) -> r.get(0));
        })
        .collectList()
        .block()
));

This still prints the wrong thing

Possible Solution

As a workaround, it seems I can use a formal batch, not the statement batch syntax (though, behind the scenes, I would imagine these should be the same thing):

System.out.println((
    Flux.from(connectionFactory.create())
        .flatMap(c -> c.createBatch()
            .add("DECLARE @t TABLE(i INT);")
            .add("INSERT INTO @t VALUES (1),(2),(3);\n")
            .add("SELECT * FROM @t;\n")
            .execute())
        .flatMap(it -> {
            System.out.println("IT: " + it);
            return it.getRowsUpdated();
        })
        .collectList()
        .block()
));

Producing:

IT: io.r2dbc.mssql.MssqlResult@134650ed
IT: io.r2dbc.mssql.MssqlResult@2463bd0a
[3, 3]

[mp911de]

The problem with variable identification is somewhat that we can’t know whether something is a variable because a lot of things in the SQL server universe start with @ (DECLARE @, SELECT @@…). I think that’s a general bug in the RPC result handling. Feel free to create another ticket to align Batch vs. non-batch result behavior. We can use this one to fix the Result behavior observed from Batch.

[mp911de]

Confirming that Batch works as expected while createStatement(…) doesn’t split up the Result:

Batch

Request: SQLBatch [sql="DECLARE @t TABLE(i INT); INSERT INTO @t VALUES (1),(2),(3); SELECT * FROM @t;"]

// first Result
Response: DoneToken [done=false, hasCount=true, rowCount=3, hasMore=true, currentCommand=0]

// second Result
Response: ColumnMetadataToken [columns=[Column [name='i", type=MutableTypeInformation [maxLength=4, lengthStrategy=BYTELENTYPE, precision=10, displaySize=11, scale=0, flags=9, serverType=int, userType=0, udtTypeName="null", collation=null, charset=null], table=null]]]
Response: io.r2dbc.mssql.message.token.RowToken@1c741110
Response: io.r2dbc.mssql.message.token.RowToken@6701ff63
Response: io.r2dbc.mssql.message.token.RowToken@3d40a7c2
Response: DoneToken [done=true, hasCount=true, rowCount=3, hasMore=false, currentCommand=0]

createStatement(…).execute()

Request: RPCRequest [procName='null', procId=10, optionFlags=io.r2dbc.mssql.message.token.RpcRequest$OptionFlags@1136b469, statusFlags=0, parameterDescriptors=[RpcString [name='null', value=DECLARE @t TABLE(i INT);INSERT INTO @t VALUES (1),(2),(3);SELECT * FROM @t;], RpcString [name='null', value=]]]

// Single Result object
Response: DoneInProcToken [done=false, hasCount=true, rowCount=3, hasMore=true, currentCommand=0]
Response: ColumnMetadataToken [columns=[Column [name='i", type=MutableTypeInformation [maxLength=4, lengthStrategy=BYTELENTYPE, precision=10, displaySize=11, scale=0, flags=9, serverType=int, userType=0, udtTypeName="null", collation=null, charset=null], table=null]]]
Response: io.r2dbc.mssql.message.token.RowToken@742ddf3a
Response: io.r2dbc.mssql.message.token.RowToken@59c66b51
Response: io.r2dbc.mssql.message.token.RowToken@e7bd10e
Response: DoneInProcToken [done=false, hasCount=true, rowCount=3, hasMore=true, currentCommand=0]
Response: ReturnStatus [status=0]
Response: DoneProcToken [done=true, hasCount=false, rowCount=0, hasMore=false, currentCommand=0]