@@ -1245,7 +1245,13 @@ public void executeStreamingSql(
12451245 throw firstRes .getException ();
12461246 case UPDATE_COUNT :
12471247 returnPartialResultSet (
1248- session , 0L , !isPartitioned , responseObserver , request .getTransaction (), false );
1248+ session ,
1249+ 0L ,
1250+ !isPartitioned ,
1251+ responseObserver ,
1252+ request .getTransaction (),
1253+ transactionId ,
1254+ false );
12491255 break ;
12501256 case RESULT_SET :
12511257 default :
@@ -1290,7 +1296,8 @@ public void executeStreamingSql(
12901296 res .getUpdateCount (),
12911297 !isPartitioned ,
12921298 responseObserver ,
1293- request .getTransaction ());
1299+ request .getTransaction (),
1300+ transactionId );
12941301 break ;
12951302 default :
12961303 throw new IllegalStateException ("Unknown result type: " + res .getType ());
@@ -1798,32 +1805,30 @@ private void returnPartialResultSet(
17981805 Long updateCount ,
17991806 boolean exact ,
18001807 StreamObserver <PartialResultSet > responseObserver ,
1801- TransactionSelector transaction ) {
1802- returnPartialResultSet (session , updateCount , exact , responseObserver , transaction , true );
1808+ TransactionSelector transactionSelector ,
1809+ ByteString transactionId ) {
1810+ returnPartialResultSet (
1811+ session , updateCount , exact , responseObserver , transactionSelector , transactionId , true );
18031812 }
18041813
18051814 private void returnPartialResultSet (
18061815 Session session ,
18071816 Long updateCount ,
18081817 boolean exact ,
18091818 StreamObserver <PartialResultSet > responseObserver ,
1810- TransactionSelector transaction ,
1819+ TransactionSelector transactionSelector ,
1820+ ByteString transactionId ,
18111821 boolean complete ) {
1812- Field field =
1813- Field .newBuilder ()
1814- .setName ("UPDATE_COUNT" )
1815- .setType (Type .newBuilder ().setCode (TypeCode .INT64 ).build ())
1816- .build ();
18171822 if (exact ) {
18181823 responseObserver .onNext (
18191824 PartialResultSet .newBuilder ()
18201825 .setMetadata (
18211826 ResultSetMetadata .newBuilder ()
1822- .setRowType (StructType .newBuilder ().addFields ( field ). build ())
1827+ .setRowType (StructType .newBuilder ().build ())
18231828 .setTransaction (
1824- ignoreInlineBeginRequest .get ()
1829+ ignoreInlineBeginRequest .get () || ! transactionSelector . hasBegin ()
18251830 ? Transaction .getDefaultInstance ()
1826- : Transaction .newBuilder ().setId (transaction . getId () ).build ())
1831+ : Transaction .newBuilder ().setId (transactionId ).build ())
18271832 .build ())
18281833 .setStats (ResultSetStats .newBuilder ().setRowCountExact (updateCount ).build ())
18291834 .build ());
@@ -1832,11 +1837,11 @@ private void returnPartialResultSet(
18321837 PartialResultSet .newBuilder ()
18331838 .setMetadata (
18341839 ResultSetMetadata .newBuilder ()
1835- .setRowType (StructType .newBuilder ().addFields ( field ). build ())
1840+ .setRowType (StructType .newBuilder ().build ())
18361841 .setTransaction (
1837- ignoreInlineBeginRequest .get ()
1842+ ignoreInlineBeginRequest .get () || ! transactionSelector . hasBegin ()
18381843 ? Transaction .getDefaultInstance ()
1839- : Transaction .newBuilder ().setId (transaction . getId () ).build ())
1844+ : Transaction .newBuilder ().setId (transactionId ).build ())
18401845 .build ())
18411846 .setStats (ResultSetStats .newBuilder ().setRowCountLowerBound (updateCount ).build ())
18421847 .build ());
0 commit comments