github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/query/result_set_test.go (about) 1 package query 2 3 import ( 4 "context" 5 "fmt" 6 "io" 7 "testing" 8 9 "github.com/stretchr/testify/require" 10 "github.com/ydb-platform/ydb-go-genproto/protos/Ydb" 11 "github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query" 12 "go.uber.org/mock/gomock" 13 grpcCodes "google.golang.org/grpc/codes" 14 grpcStatus "google.golang.org/grpc/status" 15 16 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors" 17 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest" 18 ) 19 20 func TestResultSetNext(t *testing.T) { 21 ctx := xtest.Context(t) 22 ctrl := gomock.NewController(t) 23 t.Run("Next", func(t *testing.T) { 24 t.Run("EmptyResultSet", func(t *testing.T) { 25 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 26 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 27 Status: Ydb.StatusIds_SUCCESS, 28 ResultSetIndex: 0, 29 ResultSet: &Ydb.ResultSet{ 30 Columns: []*Ydb.Column{ 31 { 32 Name: "a", 33 Type: &Ydb.Type{ 34 Type: &Ydb.Type_TypeId{ 35 TypeId: Ydb.Type_UINT64, 36 }, 37 }, 38 }, 39 { 40 Name: "b", 41 Type: &Ydb.Type{ 42 Type: &Ydb.Type_TypeId{ 43 TypeId: Ydb.Type_UTF8, 44 }, 45 }, 46 }, 47 }, 48 Rows: []*Ydb.Value{}, 49 }, 50 }, nil) 51 stream.EXPECT().Recv().Return(nil, io.EOF) 52 recv, err := stream.Recv() 53 require.NoError(t, err) 54 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 55 part, err := stream.Recv() 56 if err != nil { 57 return nil, xerrors.WithStackTrace(err) 58 } 59 60 return part, nil 61 }, recv) 62 require.EqualValues(t, 0, rs.Index()) 63 { 64 _, err := rs.nextRow(ctx) 65 require.ErrorIs(t, err, io.EOF) 66 } 67 }) 68 t.Run("SecondResultSetEmpty", func(t *testing.T) { 69 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 70 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 71 Status: Ydb.StatusIds_SUCCESS, 72 ResultSetIndex: 0, 73 ResultSet: &Ydb.ResultSet{ 74 Columns: []*Ydb.Column{ 75 { 76 Name: "a", 77 Type: &Ydb.Type{ 78 Type: &Ydb.Type_TypeId{ 79 TypeId: Ydb.Type_UINT64, 80 }, 81 }, 82 }, 83 { 84 Name: "b", 85 Type: &Ydb.Type{ 86 Type: &Ydb.Type_TypeId{ 87 TypeId: Ydb.Type_UTF8, 88 }, 89 }, 90 }, 91 }, 92 Rows: []*Ydb.Value{ 93 { 94 Items: []*Ydb.Value{{ 95 Value: &Ydb.Value_Uint64Value{ 96 Uint64Value: 1, 97 }, 98 }, { 99 Value: &Ydb.Value_TextValue{ 100 TextValue: "1", 101 }, 102 }}, 103 }, 104 { 105 Items: []*Ydb.Value{{ 106 Value: &Ydb.Value_Uint64Value{ 107 Uint64Value: 2, 108 }, 109 }, { 110 Value: &Ydb.Value_TextValue{ 111 TextValue: "2", 112 }, 113 }}, 114 }, 115 { 116 Items: []*Ydb.Value{{ 117 Value: &Ydb.Value_Uint64Value{ 118 Uint64Value: 3, 119 }, 120 }, { 121 Value: &Ydb.Value_TextValue{ 122 TextValue: "3", 123 }, 124 }}, 125 }, 126 }, 127 }, 128 }, nil) 129 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 130 Status: Ydb.StatusIds_SUCCESS, 131 ResultSetIndex: 0, 132 ResultSet: &Ydb.ResultSet{ 133 Rows: []*Ydb.Value{}, 134 }, 135 }, nil) 136 stream.EXPECT().Recv().Return(nil, io.EOF) 137 recv, err := stream.Recv() 138 require.NoError(t, err) 139 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 140 part, err := stream.Recv() 141 if err != nil { 142 return nil, xerrors.WithStackTrace(err) 143 } 144 145 return part, nil 146 }, recv) 147 require.EqualValues(t, 0, rs.Index()) 148 { 149 _, err := rs.nextRow(ctx) 150 require.NoError(t, err) 151 require.EqualValues(t, 0, rs.rowIndex) 152 } 153 { 154 _, err := rs.nextRow(ctx) 155 require.NoError(t, err) 156 require.EqualValues(t, 1, rs.rowIndex) 157 } 158 { 159 _, err := rs.nextRow(ctx) 160 require.NoError(t, err) 161 require.EqualValues(t, 2, rs.rowIndex) 162 } 163 { 164 _, err := rs.nextRow(ctx) 165 require.ErrorIs(t, err, io.EOF) 166 } 167 }) 168 t.Run("IntermediateResultSetEmpty", func(t *testing.T) { 169 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 170 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 171 Status: Ydb.StatusIds_SUCCESS, 172 ResultSetIndex: 0, 173 ResultSet: &Ydb.ResultSet{ 174 Columns: []*Ydb.Column{ 175 { 176 Name: "a", 177 Type: &Ydb.Type{ 178 Type: &Ydb.Type_TypeId{ 179 TypeId: Ydb.Type_UINT64, 180 }, 181 }, 182 }, 183 { 184 Name: "b", 185 Type: &Ydb.Type{ 186 Type: &Ydb.Type_TypeId{ 187 TypeId: Ydb.Type_UTF8, 188 }, 189 }, 190 }, 191 }, 192 Rows: []*Ydb.Value{ 193 { 194 Items: []*Ydb.Value{{ 195 Value: &Ydb.Value_Uint64Value{ 196 Uint64Value: 1, 197 }, 198 }, { 199 Value: &Ydb.Value_TextValue{ 200 TextValue: "1", 201 }, 202 }}, 203 }, 204 { 205 Items: []*Ydb.Value{{ 206 Value: &Ydb.Value_Uint64Value{ 207 Uint64Value: 2, 208 }, 209 }, { 210 Value: &Ydb.Value_TextValue{ 211 TextValue: "2", 212 }, 213 }}, 214 }, 215 { 216 Items: []*Ydb.Value{{ 217 Value: &Ydb.Value_Uint64Value{ 218 Uint64Value: 3, 219 }, 220 }, { 221 Value: &Ydb.Value_TextValue{ 222 TextValue: "3", 223 }, 224 }}, 225 }, 226 }, 227 }, 228 }, nil) 229 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 230 Status: Ydb.StatusIds_SUCCESS, 231 ResultSetIndex: 0, 232 ResultSet: &Ydb.ResultSet{ 233 Rows: []*Ydb.Value{}, 234 }, 235 }, nil) 236 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 237 Status: Ydb.StatusIds_SUCCESS, 238 ResultSetIndex: 0, 239 ResultSet: &Ydb.ResultSet{ 240 Columns: []*Ydb.Column{ 241 { 242 Name: "a", 243 Type: &Ydb.Type{ 244 Type: &Ydb.Type_TypeId{ 245 TypeId: Ydb.Type_UINT64, 246 }, 247 }, 248 }, 249 { 250 Name: "b", 251 Type: &Ydb.Type{ 252 Type: &Ydb.Type_TypeId{ 253 TypeId: Ydb.Type_UTF8, 254 }, 255 }, 256 }, 257 }, 258 Rows: []*Ydb.Value{ 259 { 260 Items: []*Ydb.Value{{ 261 Value: &Ydb.Value_Uint64Value{ 262 Uint64Value: 1, 263 }, 264 }, { 265 Value: &Ydb.Value_TextValue{ 266 TextValue: "1", 267 }, 268 }}, 269 }, 270 { 271 Items: []*Ydb.Value{{ 272 Value: &Ydb.Value_Uint64Value{ 273 Uint64Value: 2, 274 }, 275 }, { 276 Value: &Ydb.Value_TextValue{ 277 TextValue: "2", 278 }, 279 }}, 280 }, 281 { 282 Items: []*Ydb.Value{{ 283 Value: &Ydb.Value_Uint64Value{ 284 Uint64Value: 3, 285 }, 286 }, { 287 Value: &Ydb.Value_TextValue{ 288 TextValue: "3", 289 }, 290 }}, 291 }, 292 }, 293 }, 294 }, nil) 295 stream.EXPECT().Recv().Return(nil, io.EOF) 296 recv, err := stream.Recv() 297 require.NoError(t, err) 298 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 299 part, err := stream.Recv() 300 if err != nil { 301 return nil, xerrors.WithStackTrace(err) 302 } 303 304 return part, nil 305 }, recv) 306 require.EqualValues(t, 0, rs.Index()) 307 { 308 _, err := rs.nextRow(ctx) 309 require.NoError(t, err) 310 require.EqualValues(t, 0, rs.rowIndex) 311 } 312 { 313 _, err := rs.nextRow(ctx) 314 require.NoError(t, err) 315 require.EqualValues(t, 1, rs.rowIndex) 316 } 317 { 318 _, err := rs.nextRow(ctx) 319 require.NoError(t, err) 320 require.EqualValues(t, 2, rs.rowIndex) 321 } 322 { 323 _, err := rs.nextRow(ctx) 324 require.NoError(t, err) 325 require.EqualValues(t, 0, rs.rowIndex) 326 } 327 { 328 _, err := rs.nextRow(ctx) 329 require.NoError(t, err) 330 require.EqualValues(t, 1, rs.rowIndex) 331 } 332 { 333 _, err := rs.nextRow(ctx) 334 require.NoError(t, err) 335 require.EqualValues(t, 2, rs.rowIndex) 336 } 337 { 338 _, err := rs.nextRow(ctx) 339 require.ErrorIs(t, err, io.EOF) 340 } 341 }) 342 t.Run("OverTwoParts", func(t *testing.T) { 343 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 344 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 345 Status: Ydb.StatusIds_SUCCESS, 346 ResultSetIndex: 0, 347 ResultSet: &Ydb.ResultSet{ 348 Columns: []*Ydb.Column{ 349 { 350 Name: "a", 351 Type: &Ydb.Type{ 352 Type: &Ydb.Type_TypeId{ 353 TypeId: Ydb.Type_UINT64, 354 }, 355 }, 356 }, 357 { 358 Name: "b", 359 Type: &Ydb.Type{ 360 Type: &Ydb.Type_TypeId{ 361 TypeId: Ydb.Type_UTF8, 362 }, 363 }, 364 }, 365 }, 366 Rows: []*Ydb.Value{ 367 { 368 Items: []*Ydb.Value{{ 369 Value: &Ydb.Value_Uint64Value{ 370 Uint64Value: 1, 371 }, 372 }, { 373 Value: &Ydb.Value_TextValue{ 374 TextValue: "1", 375 }, 376 }}, 377 }, 378 { 379 Items: []*Ydb.Value{{ 380 Value: &Ydb.Value_Uint64Value{ 381 Uint64Value: 2, 382 }, 383 }, { 384 Value: &Ydb.Value_TextValue{ 385 TextValue: "2", 386 }, 387 }}, 388 }, 389 { 390 Items: []*Ydb.Value{{ 391 Value: &Ydb.Value_Uint64Value{ 392 Uint64Value: 3, 393 }, 394 }, { 395 Value: &Ydb.Value_TextValue{ 396 TextValue: "3", 397 }, 398 }}, 399 }, 400 }, 401 }, 402 }, nil) 403 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 404 Status: Ydb.StatusIds_SUCCESS, 405 ResultSetIndex: 0, 406 ResultSet: &Ydb.ResultSet{ 407 Rows: []*Ydb.Value{ 408 { 409 Items: []*Ydb.Value{{ 410 Value: &Ydb.Value_Uint64Value{ 411 Uint64Value: 4, 412 }, 413 }, { 414 Value: &Ydb.Value_TextValue{ 415 TextValue: "4", 416 }, 417 }}, 418 }, 419 { 420 Items: []*Ydb.Value{{ 421 Value: &Ydb.Value_Uint64Value{ 422 Uint64Value: 5, 423 }, 424 }, { 425 Value: &Ydb.Value_TextValue{ 426 TextValue: "5", 427 }, 428 }}, 429 }, 430 }, 431 }, 432 }, nil) 433 stream.EXPECT().Recv().Return(nil, io.EOF) 434 recv, err := stream.Recv() 435 require.NoError(t, err) 436 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 437 part, err := stream.Recv() 438 if err != nil { 439 return nil, xerrors.WithStackTrace(err) 440 } 441 442 return part, nil 443 }, recv) 444 require.EqualValues(t, 0, rs.Index()) 445 { 446 _, err := rs.nextRow(ctx) 447 require.NoError(t, err) 448 require.EqualValues(t, 0, rs.rowIndex) 449 } 450 { 451 _, err := rs.nextRow(ctx) 452 require.NoError(t, err) 453 require.EqualValues(t, 1, rs.rowIndex) 454 } 455 { 456 _, err := rs.nextRow(ctx) 457 require.NoError(t, err) 458 require.EqualValues(t, 2, rs.rowIndex) 459 } 460 { 461 _, err := rs.nextRow(ctx) 462 require.NoError(t, err) 463 require.EqualValues(t, 0, rs.rowIndex) 464 } 465 { 466 _, err := rs.nextRow(ctx) 467 require.NoError(t, err) 468 require.EqualValues(t, 1, rs.rowIndex) 469 } 470 { 471 _, err := rs.nextRow(ctx) 472 require.ErrorIs(t, err, io.EOF) 473 } 474 }) 475 t.Run("CanceledContext", func(t *testing.T) { 476 childCtx, cancel := context.WithCancel(xtest.Context(t)) 477 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 478 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 479 Status: Ydb.StatusIds_SUCCESS, 480 ResultSetIndex: 0, 481 ResultSet: &Ydb.ResultSet{ 482 Columns: []*Ydb.Column{ 483 { 484 Name: "a", 485 Type: &Ydb.Type{ 486 Type: &Ydb.Type_TypeId{ 487 TypeId: Ydb.Type_UINT64, 488 }, 489 }, 490 }, 491 { 492 Name: "b", 493 Type: &Ydb.Type{ 494 Type: &Ydb.Type_TypeId{ 495 TypeId: Ydb.Type_UTF8, 496 }, 497 }, 498 }, 499 }, 500 Rows: []*Ydb.Value{ 501 { 502 Items: []*Ydb.Value{{ 503 Value: &Ydb.Value_Uint64Value{ 504 Uint64Value: 1, 505 }, 506 }, { 507 Value: &Ydb.Value_TextValue{ 508 TextValue: "1", 509 }, 510 }}, 511 }, 512 { 513 Items: []*Ydb.Value{{ 514 Value: &Ydb.Value_Uint64Value{ 515 Uint64Value: 2, 516 }, 517 }, { 518 Value: &Ydb.Value_TextValue{ 519 TextValue: "2", 520 }, 521 }}, 522 }, 523 { 524 Items: []*Ydb.Value{{ 525 Value: &Ydb.Value_Uint64Value{ 526 Uint64Value: 3, 527 }, 528 }, { 529 Value: &Ydb.Value_TextValue{ 530 TextValue: "3", 531 }, 532 }}, 533 }, 534 }, 535 }, 536 }, nil) 537 recv, err := stream.Recv() 538 require.NoError(t, err) 539 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 540 part, err := stream.Recv() 541 if err != nil { 542 return nil, xerrors.WithStackTrace(err) 543 } 544 545 return part, nil 546 }, recv) 547 require.EqualValues(t, 0, rs.Index()) 548 { 549 _, err := rs.nextRow(childCtx) 550 require.NoError(t, err) 551 require.EqualValues(t, 0, rs.rowIndex) 552 } 553 cancel() 554 { 555 _, err := rs.nextRow(childCtx) 556 require.ErrorIs(t, err, context.Canceled) 557 } 558 }) 559 t.Run("OperationError", func(t *testing.T) { 560 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 561 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 562 Status: Ydb.StatusIds_SUCCESS, 563 ResultSetIndex: 0, 564 ResultSet: &Ydb.ResultSet{ 565 Columns: []*Ydb.Column{ 566 { 567 Name: "a", 568 Type: &Ydb.Type{ 569 Type: &Ydb.Type_TypeId{ 570 TypeId: Ydb.Type_UINT64, 571 }, 572 }, 573 }, 574 { 575 Name: "b", 576 Type: &Ydb.Type{ 577 Type: &Ydb.Type_TypeId{ 578 TypeId: Ydb.Type_UTF8, 579 }, 580 }, 581 }, 582 }, 583 Rows: []*Ydb.Value{ 584 { 585 Items: []*Ydb.Value{{ 586 Value: &Ydb.Value_Uint64Value{ 587 Uint64Value: 1, 588 }, 589 }, { 590 Value: &Ydb.Value_TextValue{ 591 TextValue: "1", 592 }, 593 }}, 594 }, 595 { 596 Items: []*Ydb.Value{{ 597 Value: &Ydb.Value_Uint64Value{ 598 Uint64Value: 2, 599 }, 600 }, { 601 Value: &Ydb.Value_TextValue{ 602 TextValue: "2", 603 }, 604 }}, 605 }, 606 { 607 Items: []*Ydb.Value{{ 608 Value: &Ydb.Value_Uint64Value{ 609 Uint64Value: 3, 610 }, 611 }, { 612 Value: &Ydb.Value_TextValue{ 613 TextValue: "3", 614 }, 615 }}, 616 }, 617 }, 618 }, 619 }, nil) 620 stream.EXPECT().Recv().Return(nil, xerrors.Operation(xerrors.WithStatusCode( 621 Ydb.StatusIds_OVERLOADED, 622 ))) 623 recv, err := stream.Recv() 624 require.NoError(t, err) 625 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 626 part, err := nextPart(stream) 627 if err != nil { 628 return nil, xerrors.WithStackTrace(err) 629 } 630 if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 { 631 return nil, xerrors.WithStackTrace(fmt.Errorf( 632 "critical violation of the logic: wrong result set rowIndex: %d != %d", 633 resultSetIndex, 0, 634 )) 635 } 636 637 return part, nil 638 }, recv) 639 require.EqualValues(t, 0, rs.Index()) 640 { 641 _, err := rs.nextRow(ctx) 642 require.NoError(t, err) 643 require.EqualValues(t, 0, rs.rowIndex) 644 } 645 { 646 _, err := rs.nextRow(ctx) 647 require.NoError(t, err) 648 require.EqualValues(t, 1, rs.rowIndex) 649 } 650 { 651 _, err := rs.nextRow(ctx) 652 require.NoError(t, err) 653 require.EqualValues(t, 2, rs.rowIndex) 654 } 655 { 656 _, err := rs.nextRow(ctx) 657 require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_OVERLOADED)) 658 } 659 }) 660 t.Run("TransportError", func(t *testing.T) { 661 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 662 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 663 Status: Ydb.StatusIds_SUCCESS, 664 ResultSetIndex: 0, 665 ResultSet: &Ydb.ResultSet{ 666 Columns: []*Ydb.Column{ 667 { 668 Name: "a", 669 Type: &Ydb.Type{ 670 Type: &Ydb.Type_TypeId{ 671 TypeId: Ydb.Type_UINT64, 672 }, 673 }, 674 }, 675 { 676 Name: "b", 677 Type: &Ydb.Type{ 678 Type: &Ydb.Type_TypeId{ 679 TypeId: Ydb.Type_UTF8, 680 }, 681 }, 682 }, 683 }, 684 Rows: []*Ydb.Value{ 685 { 686 Items: []*Ydb.Value{{ 687 Value: &Ydb.Value_Uint64Value{ 688 Uint64Value: 1, 689 }, 690 }, { 691 Value: &Ydb.Value_TextValue{ 692 TextValue: "1", 693 }, 694 }}, 695 }, 696 { 697 Items: []*Ydb.Value{{ 698 Value: &Ydb.Value_Uint64Value{ 699 Uint64Value: 2, 700 }, 701 }, { 702 Value: &Ydb.Value_TextValue{ 703 TextValue: "2", 704 }, 705 }}, 706 }, 707 { 708 Items: []*Ydb.Value{{ 709 Value: &Ydb.Value_Uint64Value{ 710 Uint64Value: 3, 711 }, 712 }, { 713 Value: &Ydb.Value_TextValue{ 714 TextValue: "3", 715 }, 716 }}, 717 }, 718 }, 719 }, 720 }, nil) 721 stream.EXPECT().Recv().Return(nil, grpcStatus.Error(grpcCodes.Unavailable, "")) 722 recv, err := stream.Recv() 723 require.NoError(t, err) 724 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 725 part, err := nextPart(stream) 726 if err != nil { 727 return nil, xerrors.WithStackTrace(err) 728 } 729 if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 { 730 return nil, xerrors.WithStackTrace(fmt.Errorf( 731 "critical violation of the logic: wrong result set rowIndex: %d != %d", 732 resultSetIndex, 0, 733 )) 734 } 735 736 return part, nil 737 }, recv) 738 require.EqualValues(t, 0, rs.Index()) 739 { 740 _, err := rs.nextRow(ctx) 741 require.NoError(t, err) 742 require.EqualValues(t, 0, rs.rowIndex) 743 } 744 { 745 _, err := rs.nextRow(ctx) 746 require.NoError(t, err) 747 require.EqualValues(t, 1, rs.rowIndex) 748 } 749 { 750 _, err := rs.nextRow(ctx) 751 require.NoError(t, err) 752 require.EqualValues(t, 2, rs.rowIndex) 753 } 754 { 755 _, err := rs.nextRow(ctx) 756 require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable)) 757 } 758 }) 759 t.Run("WrongResultSetIndex", func(t *testing.T) { 760 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 761 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 762 Status: Ydb.StatusIds_SUCCESS, 763 ResultSetIndex: 0, 764 ResultSet: &Ydb.ResultSet{ 765 Columns: []*Ydb.Column{ 766 { 767 Name: "a", 768 Type: &Ydb.Type{ 769 Type: &Ydb.Type_TypeId{ 770 TypeId: Ydb.Type_UINT64, 771 }, 772 }, 773 }, 774 { 775 Name: "b", 776 Type: &Ydb.Type{ 777 Type: &Ydb.Type_TypeId{ 778 TypeId: Ydb.Type_UTF8, 779 }, 780 }, 781 }, 782 }, 783 Rows: []*Ydb.Value{ 784 { 785 Items: []*Ydb.Value{{ 786 Value: &Ydb.Value_Uint64Value{ 787 Uint64Value: 1, 788 }, 789 }, { 790 Value: &Ydb.Value_TextValue{ 791 TextValue: "1", 792 }, 793 }}, 794 }, 795 { 796 Items: []*Ydb.Value{{ 797 Value: &Ydb.Value_Uint64Value{ 798 Uint64Value: 2, 799 }, 800 }, { 801 Value: &Ydb.Value_TextValue{ 802 TextValue: "2", 803 }, 804 }}, 805 }, 806 { 807 Items: []*Ydb.Value{{ 808 Value: &Ydb.Value_Uint64Value{ 809 Uint64Value: 3, 810 }, 811 }, { 812 Value: &Ydb.Value_TextValue{ 813 TextValue: "3", 814 }, 815 }}, 816 }, 817 }, 818 }, 819 }, nil) 820 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 821 Status: Ydb.StatusIds_SUCCESS, 822 ResultSetIndex: 1, 823 ResultSet: &Ydb.ResultSet{ 824 Columns: []*Ydb.Column{ 825 { 826 Name: "a", 827 Type: &Ydb.Type{ 828 Type: &Ydb.Type_TypeId{ 829 TypeId: Ydb.Type_UINT64, 830 }, 831 }, 832 }, 833 { 834 Name: "b", 835 Type: &Ydb.Type{ 836 Type: &Ydb.Type_TypeId{ 837 TypeId: Ydb.Type_UTF8, 838 }, 839 }, 840 }, 841 }, 842 Rows: []*Ydb.Value{ 843 { 844 Items: []*Ydb.Value{{ 845 Value: &Ydb.Value_Uint64Value{ 846 Uint64Value: 1, 847 }, 848 }, { 849 Value: &Ydb.Value_TextValue{ 850 TextValue: "1", 851 }, 852 }}, 853 }, 854 { 855 Items: []*Ydb.Value{{ 856 Value: &Ydb.Value_Uint64Value{ 857 Uint64Value: 2, 858 }, 859 }, { 860 Value: &Ydb.Value_TextValue{ 861 TextValue: "2", 862 }, 863 }}, 864 }, 865 { 866 Items: []*Ydb.Value{{ 867 Value: &Ydb.Value_Uint64Value{ 868 Uint64Value: 3, 869 }, 870 }, { 871 Value: &Ydb.Value_TextValue{ 872 TextValue: "3", 873 }, 874 }}, 875 }, 876 }, 877 }, 878 }, nil) 879 recv, err := stream.Recv() 880 require.NoError(t, err) 881 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 882 part, err := nextPart(stream) 883 if err != nil { 884 return nil, xerrors.WithStackTrace(err) 885 } 886 887 return part, nil 888 }, recv) 889 require.EqualValues(t, 0, rs.Index()) 890 { 891 _, err := rs.nextRow(ctx) 892 require.NoError(t, err) 893 require.EqualValues(t, 0, rs.rowIndex) 894 } 895 { 896 _, err := rs.nextRow(ctx) 897 require.NoError(t, err) 898 require.EqualValues(t, 1, rs.rowIndex) 899 } 900 { 901 _, err := rs.nextRow(ctx) 902 require.NoError(t, err) 903 require.EqualValues(t, 2, rs.rowIndex) 904 } 905 { 906 _, err := rs.nextRow(ctx) 907 require.ErrorIs(t, err, errWrongResultSetIndex) 908 } 909 { 910 _, err := rs.nextRow(ctx) 911 require.ErrorIs(t, err, io.EOF) 912 } 913 }) 914 }) 915 stream := NewMockQueryService_ExecuteQueryClient(ctrl) 916 stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{ 917 Status: Ydb.StatusIds_SUCCESS, 918 ResultSetIndex: 0, 919 ResultSet: &Ydb.ResultSet{ 920 Columns: []*Ydb.Column{ 921 { 922 Name: "a", 923 Type: &Ydb.Type{ 924 Type: &Ydb.Type_TypeId{ 925 TypeId: Ydb.Type_UINT64, 926 }, 927 }, 928 }, 929 { 930 Name: "b", 931 Type: &Ydb.Type{ 932 Type: &Ydb.Type_TypeId{ 933 TypeId: Ydb.Type_UTF8, 934 }, 935 }, 936 }, 937 }, 938 Rows: []*Ydb.Value{}, 939 }, 940 }, nil) 941 recv, err := stream.Recv() 942 require.NoError(t, err) 943 rs := newResultSet(func() (*Ydb_Query.ExecuteQueryResponsePart, error) { 944 part, err := stream.Recv() 945 if err != nil { 946 return nil, xerrors.WithStackTrace(err) 947 } 948 949 return part, nil 950 }, recv) 951 require.EqualValues(t, 0, rs.Index()) 952 t.Run("Columns", func(t *testing.T) { 953 require.EqualValues(t, []string{"a", "b"}, rs.Columns()) 954 }) 955 t.Run("ColumnNames", func(t *testing.T) { 956 var types []string 957 for _, tt := range rs.ColumnTypes() { 958 types = append(types, tt.Yql()) 959 } 960 require.EqualValues(t, []string{"Uint64", "Utf8"}, types) 961 }) 962 }