// TestTxnCoordSenderCleanupOnAborted verifies that if a txn receives a
// TransactionAbortedError, the coordinator cleans up the transaction.
func TestTxnCoordSenderCleanupOnAborted(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
// Create a transaction with intent at "a".
key := roachpb.Key("a")
txn1 := client.NewTxn(context.Background(), *s.DB)
txn1.InternalSetPriority(1)
if err := txn1.Put(key, []byte("value")); err != nil {
t.Fatal(err)
}
// Push the transaction (by writing key "a" with higher priority) to abort it.
txn2 := client.NewTxn(context.Background(), *s.DB)
txn2.InternalSetPriority(2)
if err := txn2.Put(key, []byte("value2")); err != nil {
t.Fatal(err)
}
// Now end the transaction and verify we've cleanup up, even though
// end transaction failed.
err := txn1.CommitOrCleanup()
assertTransactionAbortedError(t, err)
if err := txn2.CommitOrCleanup(); err != nil {
t.Fatal(err)
}
verifyCleanup(key, sender, s.Eng, t)
}
// TestTxnInitialTimestamp verifies that the timestamp requested
// before the Txn is created is honored.
func TestTxnInitialTimestamp(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
// Request a specific timestamp.
refTimestamp := s.Clock.Now().Add(42, 69)
txn.Proto.OrigTimestamp = refTimestamp
// Put request will create a new transaction.
key := roachpb.Key("key")
txn.InternalSetPriority(10)
txn.Proto.Isolation = enginepb.SNAPSHOT
txn.Proto.Name = "test txn"
if err := txn.Put(key, []byte("value")); err != nil {
t.Fatal(err)
}
if txn.Proto.OrigTimestamp != refTimestamp {
t.Errorf("expected txn orig ts to be %s; got %s", refTimestamp, txn.Proto.OrigTimestamp)
}
if txn.Proto.Timestamp != refTimestamp {
t.Errorf("expected txn ts to be %s; got %s", refTimestamp, txn.Proto.Timestamp)
}
}
// TestTxnCoordSenderAddIntentOnError verifies that intents are tracked if
// the transaction is, even on error.
func TestTxnCoordSenderAddIntentOnError(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
// Create a transaction with intent at "a".
key := roachpb.Key("x")
txn := client.NewTxn(context.Background(), *s.DB)
// Write so that the coordinator begins tracking this txn.
if err := txn.Put("x", "y"); err != nil {
t.Fatal(err)
}
err, ok := txn.CPut(key, []byte("x"), []byte("born to fail")).(*roachpb.ConditionFailedError)
if !ok {
t.Fatal(err)
}
sender.Lock()
txnID := *txn.Proto.ID
intentSpans, _ := roachpb.MergeSpans(sender.txns[txnID].keys)
expSpans := []roachpb.Span{{Key: key, EndKey: []byte("")}}
equal := !reflect.DeepEqual(intentSpans, expSpans)
sender.Unlock()
if err := txn.Rollback(); err != nil {
t.Fatal(err)
}
if !equal {
t.Fatalf("expected stored intents %v, got %v", expSpans, intentSpans)
}
}
// TestTxnCoordSenderGCTimeout verifies that the coordinator cleans up extant
// transactions and intents after the lastUpdateNanos exceeds the timeout.
func TestTxnCoordSenderGCTimeout(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
// Set heartbeat interval to 1ms for testing.
sender.heartbeatInterval = 1 * time.Millisecond
txn := client.NewTxn(context.Background(), *s.DB)
key := roachpb.Key("a")
if err := txn.Put(key, []byte("value")); err != nil {
t.Fatal(err)
}
// Now, advance clock past the default client timeout.
// Locking the TxnCoordSender to prevent a data race.
sender.Lock()
s.Manual.Increment(defaultClientTimeout.Nanoseconds() + 1)
sender.Unlock()
txnID := *txn.Proto.ID
util.SucceedsSoon(t, func() error {
// Locking the TxnCoordSender to prevent a data race.
sender.Lock()
_, ok := sender.txns[txnID]
sender.Unlock()
if ok {
return errors.Errorf("expected garbage collection")
}
return nil
})
verifyCleanup(key, sender, s.Eng, t)
}
// TestTxnCoordIdempotentCleanup verifies that cleanupTxnLocked is idempotent.
func TestTxnCoordIdempotentCleanup(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
ba := txn.NewBatch()
ba.Put(roachpb.Key("a"), []byte("value"))
if err := txn.Run(ba); err != nil {
t.Fatal(err)
}
sender.Lock()
// Clean up twice successively.
sender.cleanupTxnLocked(context.Background(), txn.Proto)
sender.cleanupTxnLocked(context.Background(), txn.Proto)
sender.Unlock()
// For good measure, try to commit (which cleans up once more if it
// succeeds, which it may not if the previous cleanup has already
// terminated the heartbeat goroutine)
ba = txn.NewBatch()
ba.AddRawRequest(&roachpb.EndTransactionRequest{})
err := txn.Run(ba)
if err != nil && !testutils.IsError(err, errNoState.Error()) {
t.Fatal(err)
}
}
// TestTxnCoordSenderBeginTransaction verifies that a command sent with a
// not-nil Txn with empty ID gets a new transaction initialized.
func TestTxnCoordSenderBeginTransaction(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
// Put request will create a new transaction.
key := roachpb.Key("key")
txn.InternalSetPriority(10)
txn.Proto.Isolation = enginepb.SNAPSHOT
txn.Proto.Name = "test txn"
if err := txn.Put(key, []byte("value")); err != nil {
t.Fatal(err)
}
if txn.Proto.Name != "test txn" {
t.Errorf("expected txn name to be %q; got %q", "test txn", txn.Proto.Name)
}
if txn.Proto.Priority != 10 {
t.Errorf("expected txn priority 10; got %d", txn.Proto.Priority)
}
if !bytes.Equal(txn.Proto.Key, key) {
t.Errorf("expected txn Key to match %q != %q", key, txn.Proto.Key)
}
if txn.Proto.Isolation != enginepb.SNAPSHOT {
t.Errorf("expected txn isolation to be SNAPSHOT; got %s", txn.Proto.Isolation)
}
}
// runTestFlow runs a flow with the given processors and returns the results.
// Any errors stop the current test.
func runTestFlow(
t *testing.T, srv serverutils.TestServerInterface, procs ...distsqlrun.ProcessorSpec,
) sqlbase.EncDatumRows {
kvDB := srv.KVClient().(*client.DB)
distSQLSrv := srv.DistSQLServer().(*distsqlrun.ServerImpl)
req := distsqlrun.SetupFlowRequest{
Txn: client.NewTxn(context.TODO(), *kvDB).Proto,
Flow: distsqlrun.FlowSpec{
FlowID: distsqlrun.FlowID{UUID: uuid.MakeV4()},
Processors: procs,
},
}
var rowBuf distsqlrun.RowBuffer
flow, err := distSQLSrv.SetupSyncFlow(context.TODO(), &req, &rowBuf)
if err != nil {
t.Fatal(err)
}
flow.Start(func() {})
flow.Wait()
flow.Cleanup()
if rowBuf.Err != nil {
t.Fatal(rowBuf.Err)
}
if !rowBuf.Closed {
t.Errorf("output not closed")
}
return rowBuf.Rows
}
// TestTxnCoordSenderMultipleTxns verifies correct operation with
// multiple outstanding transactions.
func TestTxnCoordSenderMultipleTxns(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn1 := client.NewTxn(context.Background(), *s.DB)
txn2 := client.NewTxn(context.Background(), *s.DB)
if err := txn1.Put(roachpb.Key("a"), []byte("value")); err != nil {
t.Fatal(err)
}
if err := txn2.Put(roachpb.Key("b"), []byte("value")); err != nil {
t.Fatal(err)
}
if len(sender.txns) != 2 {
t.Errorf("expected length of transactions map to be 2; got %d", len(sender.txns))
}
}
func backupRestoreTestSetup(
t testing.TB, numAccounts int,
) (
ctx context.Context,
tempDir string,
tc *testcluster.TestCluster,
kvDB *client.DB,
sqlDB *sqlutils.SQLRunner,
cleanup func(),
) {
ctx = context.Background()
dir, dirCleanupFn := testutils.TempDir(t, 1)
// Use ReplicationManual so we can force full replication, which is needed
// to later move the leases around.
tc = testcluster.StartTestCluster(t, backupRestoreClusterSize, base.TestClusterArgs{
ReplicationMode: base.ReplicationManual,
})
sqlDB = sqlutils.MakeSQLRunner(t, tc.Conns[0])
kvDB = tc.Server(0).KVClient().(*client.DB)
sqlDB.Exec(bankCreateDatabase)
sqlDB.Exec(bankCreateTable)
for _, insert := range bankDataInsertStmts(numAccounts) {
sqlDB.Exec(insert)
}
for _, split := range bankSplitStmts(numAccounts, backupRestoreDefaultRanges) {
sqlDB.Exec(split)
}
targets := make([]testcluster.ReplicationTarget, backupRestoreClusterSize-1)
for i := 1; i < backupRestoreClusterSize; i++ {
targets[i-1] = tc.Target(i)
}
txn := client.NewTxn(ctx, *kvDB)
rangeDescs, err := sql.AllRangeDescriptors(txn)
if err != nil {
t.Fatal(err)
}
for _, r := range rangeDescs {
if _, err := tc.AddReplicas(r.StartKey.AsRawKey(), targets...); err != nil {
t.Fatal(err)
}
}
cleanupFn := func() {
tc.Stopper().Stop()
dirCleanupFn()
}
return ctx, dir, tc, kvDB, sqlDB, cleanupFn
}
func rebalanceLeases(t testing.TB, tc *testcluster.TestCluster) {
kvDB := tc.Server(0).KVClient().(*client.DB)
txn := client.NewTxn(context.Background(), *kvDB)
rangeDescs, err := sql.AllRangeDescriptors(txn)
if err != nil {
t.Fatal(err)
}
for _, r := range rangeDescs {
target := tc.Target(int(r.RangeID) % tc.NumServers())
if err := tc.TransferRangeLease(r, target); err != nil {
t.Fatal(err)
}
}
}
// TestTxnCoordSenderKeyRanges verifies that multiple requests to same or
// overlapping key ranges causes the coordinator to keep track only of
// the minimum number of ranges.
func TestTxnCoordSenderKeyRanges(t *testing.T) {
defer leaktest.AfterTest(t)()
ranges := []struct {
start, end roachpb.Key
}{
{roachpb.Key("a"), roachpb.Key(nil)},
{roachpb.Key("a"), roachpb.Key(nil)},
{roachpb.Key("aa"), roachpb.Key(nil)},
{roachpb.Key("b"), roachpb.Key(nil)},
{roachpb.Key("aa"), roachpb.Key("c")},
{roachpb.Key("b"), roachpb.Key("c")},
}
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
for _, rng := range ranges {
if rng.end != nil {
if err := txn.DelRange(rng.start, rng.end); err != nil {
t.Fatal(err)
}
} else {
if err := txn.Put(rng.start, []byte("value")); err != nil {
t.Fatal(err)
}
}
}
txnID := *txn.Proto.ID
// Verify that the transaction metadata contains only two entries
// in its "keys" range group. "a" and range "aa"-"c".
txnMeta, ok := sender.txns[txnID]
if !ok {
t.Fatalf("expected a transaction to be created on coordinator")
}
keys, _ := roachpb.MergeSpans(txnMeta.keys)
if len(keys) != 2 {
t.Errorf("expected 2 entries in keys range group; got %v", keys)
}
}
// TestTxnCoordSenderReleaseTxnMeta verifies that TxnCoordSender releases the
// txnMetadata after the txn has committed successfully.
func TestTxnCoordSenderReleaseTxnMeta(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
ba := txn.NewBatch()
ba.Put(roachpb.Key("a"), []byte("value"))
ba.Put(roachpb.Key("b"), []byte("value"))
if err := txn.CommitInBatch(ba); err != nil {
t.Fatal(err)
}
txnID := *txn.Proto.ID
if _, ok := sender.txns[txnID]; ok {
t.Fatal("expected TxnCoordSender has released the txn")
}
}
// TestTxnCoordSenderBeginTransactionMinPriority verifies that when starting
// a new transaction, a non-zero priority is treated as a minimum value.
func TestTxnCoordSenderBeginTransactionMinPriority(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
// Put request will create a new transaction.
key := roachpb.Key("key")
txn.InternalSetPriority(10)
txn.Proto.Isolation = enginepb.SNAPSHOT
txn.Proto.Priority = 11
if err := txn.Put(key, []byte("value")); err != nil {
t.Fatal(err)
}
if prio := txn.Proto.Priority; prio != 11 {
t.Errorf("expected txn priority 11; got %d", prio)
}
}
// TestReadOnlyTxnObeysDeadline tests that read-only transactions obey the
// deadline.
func TestReadOnlyTxnObeysDeadline(t *testing.T) {
defer leaktest.AfterTest(t)()
s, _, _ := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop()
db := createTestClient(t, s)
if err := db.Put(context.TODO(), "k", "v"); err != nil {
t.Fatal(err)
}
// Use txn.Exec instead of db.Txn to disable auto retry.
txn := client.NewTxn(context.TODO(), *db)
if err := txn.Exec(client.TxnExecOptions{AutoRetry: false, AutoCommit: true}, func(txn *client.Txn, _ *client.TxnExecOptions) error {
// Set deadline to sometime in the past.
txn.UpdateDeadlineMaybe(hlc.Timestamp{WallTime: timeutil.Now().Add(-time.Second).UnixNano()})
_, err := txn.Get("k")
return err
}); !testutils.IsError(err, "txn aborted") {
t.Fatal(err)
}
}
// TestTxnCoordSenderAddRequest verifies adding a request creates a
// transaction metadata and adding multiple requests with same
// transaction ID updates the last update timestamp.
func TestTxnCoordSenderAddRequest(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
txn := client.NewTxn(context.Background(), *s.DB)
// Put request will create a new transaction.
if err := txn.Put(roachpb.Key("a"), []byte("value")); err != nil {
t.Fatal(err)
}
txnID := *txn.Proto.ID
txnMeta, ok := sender.txns[txnID]
if !ok {
t.Fatal("expected a transaction to be created on coordinator")
}
if !txn.Proto.Writing {
t.Fatal("txn is not marked as writing")
}
ts := txnMeta.getLastUpdate()
// Advance time and send another put request. Lock the coordinator
// to prevent a data race.
sender.Lock()
s.Manual.Increment(1)
sender.Unlock()
if err := txn.Put(roachpb.Key("a"), []byte("value")); err != nil {
t.Fatal(err)
}
if len(sender.txns) != 1 {
t.Errorf("expected length of transactions map to be 1; got %d", len(sender.txns))
}
txnMeta = sender.txns[txnID]
if lu := txnMeta.getLastUpdate(); ts >= lu {
t.Errorf("expected last update time to advance past %d; got %d", ts, lu)
} else if un := s.Manual.UnixNano(); lu != un {
t.Errorf("expected last update time to equal %d; got %d", un, lu)
}
}
func TestTxnCoordSenderCancel(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
ctx, cancel := context.WithCancel(context.Background())
origSender := sender.wrapped
sender.wrapped = client.SenderFunc(
func(ctx context.Context, args roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
if _, hasET := args.GetArg(roachpb.EndTransaction); hasET {
// Cancel the transaction while also sending it along. This tickled a
// data race in TxnCoordSender.tryAsyncAbort. See #7726.
cancel()
}
return origSender.Send(ctx, args)
})
// Create a transaction with bunch of intents.
txn := client.NewTxn(ctx, *s.DB)
batch := txn.NewBatch()
for i := 0; i < 100; i++ {
key := roachpb.Key(fmt.Sprintf("%d", i))
batch.Put(key, []byte("value"))
}
if err := txn.Run(batch); err != nil {
t.Fatal(err)
}
// Commit the transaction. Note that we cancel the transaction when the
// commit is sent which stresses the TxnCoordSender.tryAsyncAbort code
// path. We'll either succeed, get a "does not exist" error, or get a
// context canceled error. Anything else is unexpected.
err := txn.CommitOrCleanup()
if err != nil && err.Error() != context.Canceled.Error() &&
!testutils.IsError(err, "does not exist") {
t.Fatal(err)
}
}
func TestTxnRestartCount(t *testing.T) {
defer leaktest.AfterTest(t)()
_, sender, cleanupFn := setupMetricsTest(t)
defer cleanupFn()
key := []byte("key-restart")
value := []byte("value")
db := client.NewDB(sender)
// Start a transaction and do a GET. This forces a timestamp to be chosen for the transaction.
txn := client.NewTxn(context.Background(), *db)
if _, err := txn.Get(key); err != nil {
t.Fatal(err)
}
// Outside of the transaction, read the same key as was read within the transaction. This
// means that future attempts to write will increase the timestamp.
if _, err := db.Get(context.TODO(), key); err != nil {
t.Fatal(err)
}
// This put will lay down an intent, txn timestamp will increase beyond original.
if err := txn.Put(key, value); err != nil {
t.Fatal(err)
}
if !txn.Proto.OrigTimestamp.Less(txn.Proto.Timestamp) {
t.Errorf("expected timestamp to increase: %s", txn.Proto)
}
// Commit (should cause restart metric to increase).
err := txn.CommitOrCleanup()
assertTransactionRetryError(t, err)
teardownHeartbeats(sender)
checkTxnMetrics(t, sender, "restart txn", 0, 0, 0, 1, 1)
}
func TestClusterFlow(t *testing.T) {
defer leaktest.AfterTest(t)()
const numRows = 100
args := base.TestClusterArgs{ReplicationMode: base.ReplicationManual}
tc := serverutils.StartTestCluster(t, 3, args)
defer tc.Stopper().Stop()
sumDigitsFn := func(row int) parser.Datum {
sum := 0
for row > 0 {
sum += row % 10
row /= 10
}
return parser.NewDInt(parser.DInt(sum))
}
sqlutils.CreateTable(t, tc.ServerConn(0), "t",
"num INT PRIMARY KEY, digitsum INT, numstr STRING, INDEX s (digitsum)",
numRows,
sqlutils.ToRowFn(sqlutils.RowIdxFn, sumDigitsFn, sqlutils.RowEnglishFn))
kvDB := tc.Server(0).KVClient().(*client.DB)
desc := sqlbase.GetTableDescriptor(kvDB, "test", "t")
makeIndexSpan := func(start, end int) TableReaderSpan {
var span roachpb.Span
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(desc, desc.Indexes[0].ID))
span.Key = append(prefix, encoding.EncodeVarintAscending(nil, int64(start))...)
span.EndKey = append(span.EndKey, prefix...)
span.EndKey = append(span.EndKey, encoding.EncodeVarintAscending(nil, int64(end))...)
return TableReaderSpan{Span: span}
}
// Set up table readers on three hosts feeding data into a join reader on
// the third host. This is a basic test for the distributed flow
// infrastructure, including local and remote streams.
//
// Note that the ranges won't necessarily be local to the table readers, but
// that doesn't matter for the purposes of this test.
// Start a span (useful to look at spans using Lighstep).
sp, err := tracing.JoinOrNew(tracing.NewTracer(), nil, "cluster test")
if err != nil {
t.Fatal(err)
}
ctx := opentracing.ContextWithSpan(context.Background(), sp)
defer sp.Finish()
tr1 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(0, 8)},
}
tr2 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(8, 12)},
}
tr3 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(12, 100)},
}
jr := JoinReaderSpec{
Table: *desc,
OutputColumns: []uint32{2},
}
txn := client.NewTxn(ctx, *kvDB)
fid := FlowID{uuid.MakeV4()}
req1 := &SetupFlowRequest{Txn: txn.Proto}
req1.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &tr1},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{StreamID: 0, Mailbox: &MailboxSpec{TargetAddr: tc.Server(2).ServingAddr()}},
},
}},
}},
}
req2 := &SetupFlowRequest{Txn: txn.Proto}
req2.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &tr2},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{StreamID: 1, Mailbox: &MailboxSpec{TargetAddr: tc.Server(2).ServingAddr()}},
},
}},
}},
}
req3 := &SetupFlowRequest{Txn: txn.Proto}
req3.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{
{
Core: ProcessorCoreUnion{TableReader: &tr3},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{StreamID: StreamID(2)},
},
}},
},
{
Input: []InputSyncSpec{{
Type: InputSyncSpec_ORDERED,
Ordering: Ordering{Columns: []Ordering_Column{{1, Ordering_Column_ASC}}},
Streams: []StreamEndpointSpec{
{StreamID: 0, Mailbox: &MailboxSpec{}},
{StreamID: 1, Mailbox: &MailboxSpec{}},
{StreamID: StreamID(2)},
},
}},
Core: ProcessorCoreUnion{JoinReader: &jr},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{{Mailbox: &MailboxSpec{SimpleResponse: true}}},
}}},
},
}
if err := SetFlowRequestTrace(ctx, req1); err != nil {
t.Fatal(err)
}
if err := SetFlowRequestTrace(ctx, req2); err != nil {
t.Fatal(err)
}
if err := SetFlowRequestTrace(ctx, req3); err != nil {
t.Fatal(err)
}
var clients []DistSQLClient
for i := 0; i < 3; i++ {
s := tc.Server(i)
conn, err := s.RPCContext().GRPCDial(s.ServingAddr())
if err != nil {
t.Fatal(err)
}
clients = append(clients, NewDistSQLClient(conn))
}
if log.V(1) {
log.Infof(ctx, "Setting up flow on 0")
}
if resp, err := clients[0].SetupFlow(ctx, req1); err != nil {
t.Fatal(err)
} else if resp.Error != nil {
t.Fatal(resp.Error)
}
if log.V(1) {
log.Infof(ctx, "Setting up flow on 1")
}
if resp, err := clients[1].SetupFlow(ctx, req2); err != nil {
t.Fatal(err)
} else if resp.Error != nil {
t.Fatal(resp.Error)
}
if log.V(1) {
log.Infof(ctx, "Running flow on 2")
}
stream, err := clients[2].RunSimpleFlow(ctx, req3)
if err != nil {
t.Fatal(err)
}
var decoder StreamDecoder
var rows sqlbase.EncDatumRows
for {
msg, err := stream.Recv()
if err != nil {
if err == io.EOF {
break
}
t.Fatal(err)
}
err = decoder.AddMessage(msg)
if err != nil {
t.Fatal(err)
}
rows = testGetDecodedRows(t, &decoder, rows)
}
if done, trailerErr := decoder.IsDone(); !done {
t.Fatal("stream not done")
} else if trailerErr != nil {
t.Fatal("error in the stream trailer:", trailerErr)
}
// The result should be all the numbers in string form, ordered by the
// digit sum (and then by number).
var results []string
for sum := 1; sum <= 50; sum++ {
for i := 1; i <= numRows; i++ {
if int(*sumDigitsFn(i).(*parser.DInt)) == sum {
results = append(results, fmt.Sprintf("['%s']", sqlutils.IntToEnglish(i)))
}
}
}
expected := strings.Join(results, " ")
expected = "[" + expected + "]"
if rowStr := rows.String(); rowStr != expected {
t.Errorf("Result: %s\n Expected: %s\n", rowStr, expected)
}
}
// TestTxnCoordSenderTxnUpdatedOnError verifies that errors adjust the
// response transaction's timestamp and priority as appropriate.
func TestTxnCoordSenderTxnUpdatedOnError(t *testing.T) {
defer leaktest.AfterTest(t)()
origTS := makeTS(123, 0)
plus10 := origTS.Add(10, 10)
plus20 := plus10.Add(10, 0)
testCases := []struct {
pErr *roachpb.Error
expEpoch uint32
expPri int32
expTS, expOrigTS hlc.Timestamp
nodeSeen bool
}{
{
// No error, so nothing interesting either.
pErr: nil,
expEpoch: 0,
expPri: 1,
expTS: origTS,
expOrigTS: origTS,
},
{
// On uncertainty error, new epoch begins and node is seen.
// Timestamp moves ahead of the existing write.
pErr: func() *roachpb.Error {
pErr := roachpb.NewErrorWithTxn(
roachpb.NewReadWithinUncertaintyIntervalError(hlc.ZeroTimestamp, hlc.ZeroTimestamp),
&roachpb.Transaction{})
const nodeID = 1
pErr.GetTxn().UpdateObservedTimestamp(nodeID, plus10)
pErr.OriginNode = nodeID
return pErr
}(),
expEpoch: 1,
expPri: 1,
expTS: plus10,
expOrigTS: plus10,
nodeSeen: true,
},
{
// On abort, nothing changes but we get a new priority to use for
// the next attempt.
pErr: roachpb.NewErrorWithTxn(&roachpb.TransactionAbortedError{},
&roachpb.Transaction{
TxnMeta: enginepb.TxnMeta{Timestamp: plus20, Priority: 10},
}),
expPri: 10,
},
{
// On failed push, new epoch begins just past the pushed timestamp.
// Additionally, priority ratchets up to just below the pusher's.
pErr: roachpb.NewErrorWithTxn(&roachpb.TransactionPushError{
PusheeTxn: roachpb.Transaction{
TxnMeta: enginepb.TxnMeta{Timestamp: plus10, Priority: int32(10)},
},
},
&roachpb.Transaction{}),
expEpoch: 1,
expPri: 9,
expTS: plus10,
expOrigTS: plus10,
},
{
// On retry, restart with new epoch, timestamp and priority.
pErr: roachpb.NewErrorWithTxn(&roachpb.TransactionRetryError{},
&roachpb.Transaction{
TxnMeta: enginepb.TxnMeta{Timestamp: plus10, Priority: int32(10)},
},
),
expEpoch: 1,
expPri: 10,
expTS: plus10,
expOrigTS: plus10,
},
}
for i, test := range testCases {
stopper := stop.NewStopper()
manual := hlc.NewManualClock(origTS.WallTime)
clock := hlc.NewClock(manual.UnixNano, 20*time.Nanosecond)
senderFunc := func(_ context.Context, ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
var reply *roachpb.BatchResponse
if test.pErr == nil {
reply = ba.CreateReply()
}
return reply, test.pErr
}
ambient := log.AmbientContext{Tracer: tracing.NewTracer()}
ts := NewTxnCoordSender(
ambient,
senderFn(senderFunc),
clock,
false,
stopper,
MakeTxnMetrics(metric.TestSampleInterval),
)
db := client.NewDB(ts)
txn := client.NewTxn(context.Background(), *db)
txn.InternalSetPriority(1)
txn.Proto.Name = "test txn"
key := roachpb.Key("test-key")
_, err := txn.Get(key)
teardownHeartbeats(ts)
stopper.Stop()
if test.pErr != nil && err == nil {
t.Fatalf("expected an error")
}
if txn.Proto.Epoch != test.expEpoch {
t.Errorf("%d: expected epoch = %d; got %d",
i, test.expEpoch, txn.Proto.Epoch)
}
if txn.Proto.Priority != test.expPri {
t.Errorf("%d: expected priority = %d; got %d",
i, test.expPri, txn.Proto.Priority)
}
if !txn.Proto.Timestamp.Equal(test.expTS) {
t.Errorf("%d: expected timestamp to be %s; got %s",
i, test.expTS, txn.Proto.Timestamp)
}
if !txn.Proto.OrigTimestamp.Equal(test.expOrigTS) {
t.Errorf("%d: expected orig timestamp to be %s; got %s",
i, test.expOrigTS, txn.Proto.OrigTimestamp)
}
if ns := txn.Proto.ObservedTimestamps; (len(ns) != 0) != test.nodeSeen {
t.Errorf("%d: expected nodeSeen=%t, but list of hosts is %v",
i, test.nodeSeen, ns)
}
}
}
// TestTxnCoordSenderGCWithCancel verifies that the coordinator cleans up extant
// transactions and intents after transaction context is cancelled.
func TestTxnCoordSenderGCWithCancel(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
// Set heartbeat interval to 1ms for testing.
sender.heartbeatInterval = 1 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
txn := client.NewTxn(ctx, *s.DB)
key := roachpb.Key("a")
if pErr := txn.Put(key, []byte("value")); pErr != nil {
t.Fatal(pErr)
}
// Now, advance clock past the default client timeout.
// Locking the TxnCoordSender to prevent a data race.
sender.Lock()
s.Manual.Increment(defaultClientTimeout.Nanoseconds() + 1)
sender.Unlock()
txnID := *txn.Proto.ID
// Verify that the transaction is alive despite the timeout having been
// exceeded.
errStillActive := errors.New("transaction is still active")
// TODO(dan): Figure out how to run the heartbeat manually instead of this.
if err := util.RetryForDuration(1*time.Second, func() error {
// Locking the TxnCoordSender to prevent a data race.
sender.Lock()
_, ok := sender.txns[txnID]
sender.Unlock()
if !ok {
return nil
}
meta := &enginepb.MVCCMetadata{}
ok, _, _, err := s.Eng.GetProto(engine.MakeMVCCMetadataKey(key), meta)
if err != nil {
t.Fatalf("error getting MVCC metadata: %s", err)
}
if !ok || meta.Txn == nil {
return nil
}
return errStillActive
}); err != errStillActive {
t.Fatalf("expected transaction to be active, got: %v", err)
}
// After the context is cancelled, the transaction should be cleaned up.
cancel()
util.SucceedsSoon(t, func() error {
// Locking the TxnCoordSender to prevent a data race.
sender.Lock()
_, ok := sender.txns[txnID]
sender.Unlock()
if ok {
return errors.Errorf("expected garbage collection")
}
return nil
})
verifyCleanup(key, sender, s.Eng, t)
}
func TestTableReader(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop()
// Create a table where each row is:
//
// | a | b | sum | s |
// |-----------------------------------------------------------------|
// | rowId/10 | rowId%10 | rowId/10 + rowId%10 | IntToEnglish(rowId) |
aFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row / 10))
}
bFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row % 10))
}
sumFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row/10 + row%10))
}
sqlutils.CreateTable(t, sqlDB, "t",
"a INT, b INT, sum INT, s STRING, PRIMARY KEY (a,b), INDEX bs (b,s)",
99,
sqlutils.ToRowFn(aFn, bFn, sumFn, sqlutils.RowEnglishFn))
td := sqlbase.GetTableDescriptor(kvDB, "test", "t")
makeIndexSpan := func(start, end int) TableReaderSpan {
var span roachpb.Span
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(td, td.Indexes[0].ID))
span.Key = append(prefix, encoding.EncodeVarintAscending(nil, int64(start))...)
span.EndKey = append(span.EndKey, prefix...)
span.EndKey = append(span.EndKey, encoding.EncodeVarintAscending(nil, int64(end))...)
return TableReaderSpan{Span: span}
}
testCases := []struct {
spec TableReaderSpec
expected string
}{
{
spec: TableReaderSpec{
Filter: Expression{Expr: "@3 < 5 AND @2 != 3"}, // sum < 5 && b != 3
OutputColumns: []uint32{0, 1},
},
expected: "[[0 1] [0 2] [0 4] [1 0] [1 1] [1 2] [2 0] [2 1] [2 2] [3 0] [3 1] [4 0]]",
},
{
spec: TableReaderSpec{
Filter: Expression{Expr: "@3 < 5 AND @2 != 3"},
OutputColumns: []uint32{3}, // s
HardLimit: 4,
},
expected: "[['one'] ['two'] ['four'] ['one-zero']]",
},
{
spec: TableReaderSpec{
IndexIdx: 1,
Reverse: true,
Spans: []TableReaderSpan{makeIndexSpan(4, 6)},
Filter: Expression{Expr: "@1 < 3"}, // sum < 8
OutputColumns: []uint32{0, 1},
SoftLimit: 1,
},
expected: "[[2 5] [1 5] [0 5] [2 4] [1 4] [0 4]]",
},
}
for _, c := range testCases {
ts := c.spec
ts.Table = *td
txn := client.NewTxn(context.Background(), *kvDB)
flowCtx := FlowCtx{
Context: context.Background(),
evalCtx: &parser.EvalContext{},
txn: txn,
}
out := &RowBuffer{}
tr, err := newTableReader(&flowCtx, &ts, out)
if err != nil {
t.Fatal(err)
}
tr.Run(nil)
if out.err != nil {
t.Fatal(out.err)
}
if !out.closed {
t.Fatalf("output RowReceiver not closed")
}
if result := out.rows.String(); result != c.expected {
t.Errorf("invalid results: %s, expected %s'", result, c.expected)
}
}
}
// TestTxnCoordSenderNoDuplicateIntents verifies that TxnCoordSender does not
// generate duplicate intents and that it merges intents for overlapping ranges.
func TestTxnCoordSenderNoDuplicateIntents(t *testing.T) {
defer leaktest.AfterTest(t)()
stopper := stop.NewStopper()
manual := hlc.NewManualClock(123)
clock := hlc.NewClock(manual.UnixNano, time.Nanosecond)
var expectedIntents []roachpb.Span
senderFunc := func(_ context.Context, ba roachpb.BatchRequest) (
*roachpb.BatchResponse, *roachpb.Error) {
if rArgs, ok := ba.GetArg(roachpb.EndTransaction); ok {
et := rArgs.(*roachpb.EndTransactionRequest)
if !reflect.DeepEqual(et.IntentSpans, expectedIntents) {
t.Errorf("Invalid intents: %+v; expected %+v", et.IntentSpans, expectedIntents)
}
}
br := ba.CreateReply()
txnClone := ba.Txn.Clone()
br.Txn = &txnClone
br.Txn.Writing = true
return br, nil
}
ambient := log.AmbientContext{Tracer: tracing.NewTracer()}
ts := NewTxnCoordSender(
ambient,
senderFn(senderFunc),
clock,
false,
stopper,
MakeTxnMetrics(metric.TestSampleInterval),
)
defer stopper.Stop()
defer teardownHeartbeats(ts)
db := client.NewDB(ts)
txn := client.NewTxn(context.Background(), *db)
// Write to a, b, u-w before the final batch.
pErr := txn.Put(roachpb.Key("a"), []byte("value"))
if pErr != nil {
t.Fatal(pErr)
}
pErr = txn.Put(roachpb.Key("b"), []byte("value"))
if pErr != nil {
t.Fatal(pErr)
}
pErr = txn.DelRange(roachpb.Key("u"), roachpb.Key("w"))
if pErr != nil {
t.Fatal(pErr)
}
// The final batch overwrites key a and overlaps part of the u-w range.
b := txn.NewBatch()
b.Put(roachpb.Key("b"), []byte("value"))
b.Put(roachpb.Key("c"), []byte("value"))
b.DelRange(roachpb.Key("v"), roachpb.Key("z"), false)
// The expected intents are a, b, c, and u-z.
expectedIntents = []roachpb.Span{
{Key: roachpb.Key("a"), EndKey: nil},
{Key: roachpb.Key("b"), EndKey: nil},
{Key: roachpb.Key("c"), EndKey: nil},
{Key: roachpb.Key("u"), EndKey: roachpb.Key("z")},
}
pErr = txn.CommitInBatch(b)
if pErr != nil {
t.Fatal(pErr)
}
}
// Backup exports a snapshot of every kv entry into ranged sstables.
//
// The output is an sstable per range with files in the following locations:
// - /<base>/<node_id>/<key_range>/data.sst
// - <base> is given by the user and is expected to eventually be cloud storage
// - The <key_range>s are non-overlapping.
//
// TODO(dan): Bikeshed this directory structure and naming.
func Backup(
ctx context.Context, db client.DB, base string, endTime hlc.Timestamp,
) (desc sqlbase.BackupDescriptor, retErr error) {
// TODO(dan): Optionally take a start time for an incremental backup.
// TODO(dan): Take a uri for the path prefix and support various cloud storages.
// TODO(dan): Figure out how permissions should work. #6713 is tracking this
// for grpc.
var rangeDescs []roachpb.RangeDescriptor
var sqlDescs []sqlbase.Descriptor
opt := client.TxnExecOptions{
AutoRetry: true,
AutoCommit: true,
}
{
// TODO(dan): Pick an appropriate end time and set it in the txn.
txn := client.NewTxn(ctx, db)
err := txn.Exec(opt, func(txn *client.Txn, opt *client.TxnExecOptions) error {
var err error
SetTxnTimestamps(txn, endTime)
rangeDescs, err = AllRangeDescriptors(txn)
if err != nil {
return err
}
sqlDescs, err = allSQLDescriptors(txn)
return err
})
if err != nil {
return sqlbase.BackupDescriptor{}, err
}
}
var dataSize int64
backupDescs := make([]sqlbase.BackupRangeDescriptor, len(rangeDescs))
crc := crc32.New(crc32.MakeTable(crc32.Castagnoli))
for i, rangeDesc := range rangeDescs {
backupDescs[i] = sqlbase.BackupRangeDescriptor{
StartKey: rangeDesc.StartKey.AsRawKey(),
EndKey: rangeDesc.EndKey.AsRawKey(),
StartTime: hlc.Timestamp{},
}
if backupDescs[i].StartKey.Compare(keys.LocalMax) < 0 {
backupDescs[i].StartKey = keys.LocalMax
}
nodeID := 0
dir := filepath.Join(base, fmt.Sprintf("%03d", nodeID))
dir = filepath.Join(dir, fmt.Sprintf("%x-%x", rangeDesc.StartKey, rangeDesc.EndKey))
if err := os.MkdirAll(dir, 0700); err != nil {
return sqlbase.BackupDescriptor{}, err
}
var kvs []client.KeyValue
txn := client.NewTxn(ctx, db)
err := txn.Exec(opt, func(txn *client.Txn, opt *client.TxnExecOptions) error {
var err error
SetTxnTimestamps(txn, endTime)
// TODO(dan): Iterate with some batch size.
kvs, err = txn.Scan(backupDescs[i].StartKey, backupDescs[i].EndKey, 0)
return err
})
if err != nil {
return sqlbase.BackupDescriptor{}, err
}
if len(kvs) == 0 {
if log.V(1) {
log.Infof(ctx, "skipping backup of empty range %s-%s",
backupDescs[i].StartKey, backupDescs[i].EndKey)
}
continue
}
backupDescs[i].Path = filepath.Join(dir, dataSSTableName)
writeSST := func() (writeSSTErr error) {
// This is a function so the defered Close (and resultant flush) is
// called before the checksum is computed.
sst := engine.MakeRocksDBSstFileWriter()
if err := sst.Open(backupDescs[i].Path); err != nil {
return err
}
defer func() {
if closeErr := sst.Close(); closeErr != nil && writeSSTErr == nil {
writeSSTErr = closeErr
}
}()
// TODO(dan): Move all this iteration into cpp to avoid the cgo calls.
for _, kv := range kvs {
mvccKV := engine.MVCCKeyValue{
Key: engine.MVCCKey{Key: kv.Key, Timestamp: kv.Value.Timestamp},
Value: kv.Value.RawBytes,
}
if err := sst.Add(mvccKV); err != nil {
return err
}
}
dataSize += sst.DataSize
return nil
}
if err := writeSST(); err != nil {
return sqlbase.BackupDescriptor{}, err
}
crc.Reset()
f, err := os.Open(backupDescs[i].Path)
if err != nil {
return sqlbase.BackupDescriptor{}, err
}
defer f.Close()
if _, err := io.Copy(crc, f); err != nil {
return sqlbase.BackupDescriptor{}, err
}
backupDescs[i].CRC = crc.Sum32()
}
desc = sqlbase.BackupDescriptor{
EndTime: endTime,
Ranges: backupDescs,
SQL: sqlDescs,
DataSize: dataSize,
}
descBuf, err := desc.Marshal()
if err != nil {
return sqlbase.BackupDescriptor{}, err
}
if err = ioutil.WriteFile(filepath.Join(base, backupDescriptorName), descBuf, 0600); err != nil {
return sqlbase.BackupDescriptor{}, err
}
return desc, nil
}
// TestTxnCoordSenderEndTxn verifies that ending a transaction
// sends resolve write intent requests and removes the transaction
// from the txns map.
func TestTxnCoordSenderEndTxn(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
// 4 cases: no deadline, past deadline, equal deadline, future deadline.
for i := 0; i < 4; i++ {
key := roachpb.Key("key: " + strconv.Itoa(i))
txn := client.NewTxn(context.Background(), *s.DB)
// Set to SNAPSHOT so that it can be pushed without restarting.
if err := txn.SetIsolation(enginepb.SNAPSHOT); err != nil {
t.Fatal(err)
}
// Initialize the transaction.
if pErr := txn.Put(key, []byte("value")); pErr != nil {
t.Fatal(pErr)
}
// Conflicting transaction that pushes the above transaction.
conflictTxn := client.NewTxn(context.Background(), *s.DB)
if _, pErr := conflictTxn.Get(key); pErr != nil {
t.Fatal(pErr)
}
// The transaction was pushed at least to conflictTxn's timestamp (but
// it could have been pushed more - the push takes a timestamp off the
// HLC).
pusheeTxn, pErr := getTxn(sender, &txn.Proto)
if pErr != nil {
t.Fatal(pErr)
}
pushedTimestamp := pusheeTxn.Timestamp
{
var err error
switch i {
case 0:
// No deadline.
case 1:
// Past deadline.
if !txn.UpdateDeadlineMaybe(pushedTimestamp.Prev()) {
t.Fatalf("did not update deadline")
}
case 2:
// Equal deadline.
if !txn.UpdateDeadlineMaybe(pushedTimestamp) {
t.Fatalf("did not update deadline")
}
case 3:
// Future deadline.
if !txn.UpdateDeadlineMaybe(pushedTimestamp.Next()) {
t.Fatalf("did not update deadline")
}
}
err = txn.CommitOrCleanup()
switch i {
case 0:
// No deadline.
if err != nil {
t.Fatal(err)
}
case 1:
// Past deadline.
if statusError, ok := err.(*roachpb.TransactionStatusError); !ok {
t.Fatalf("expected TransactionStatusError but got %T: %s", err, err)
} else if expected := "transaction deadline exceeded"; statusError.Msg != expected {
t.Fatalf("expected %s, got %s", expected, statusError.Msg)
}
case 2:
// Equal deadline.
if err != nil {
t.Fatal(err)
}
case 3:
// Future deadline.
if err != nil {
t.Fatal(err)
}
}
}
verifyCleanup(key, sender, s.Eng, t)
}
}
// TestTxnCoordSenderHeartbeat verifies periodic heartbeat of the
// transaction record.
func TestTxnCoordSenderHeartbeat(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sender := createTestDB(t)
defer s.Stop()
defer teardownHeartbeats(sender)
// Set heartbeat interval to 1ms for testing.
sender.heartbeatInterval = 1 * time.Millisecond
initialTxn := client.NewTxn(context.Background(), *s.DB)
if err := initialTxn.Put(roachpb.Key("a"), []byte("value")); err != nil {
t.Fatal(err)
}
// Verify 3 heartbeats.
var heartbeatTS hlc.Timestamp
for i := 0; i < 3; i++ {
util.SucceedsSoon(t, func() error {
txn, pErr := getTxn(sender, &initialTxn.Proto)
if pErr != nil {
t.Fatal(pErr)
}
// Advance clock by 1ns.
// Locking the TxnCoordSender to prevent a data race.
sender.Lock()
s.Manual.Increment(1)
sender.Unlock()
if txn.LastHeartbeat != nil && heartbeatTS.Less(*txn.LastHeartbeat) {
heartbeatTS = *txn.LastHeartbeat
return nil
}
return errors.Errorf("expected heartbeat")
})
}
// Sneakily send an ABORT right to DistSender (bypassing TxnCoordSender).
{
var ba roachpb.BatchRequest
ba.Add(&roachpb.EndTransactionRequest{
Commit: false,
Span: roachpb.Span{Key: initialTxn.Proto.Key},
})
ba.Txn = &initialTxn.Proto
if _, pErr := sender.wrapped.Send(context.Background(), ba); pErr != nil {
t.Fatal(pErr)
}
}
util.SucceedsSoon(t, func() error {
sender.Lock()
defer sender.Unlock()
if txnMeta, ok := sender.txns[*initialTxn.Proto.ID]; !ok {
t.Fatal("transaction unregistered prematurely")
} else if txnMeta.txn.Status != roachpb.ABORTED {
return fmt.Errorf("transaction is not aborted")
}
return nil
})
// Trying to do something else should give us a TransactionAbortedError.
_, err := initialTxn.Get("a")
assertTransactionAbortedError(t, err)
}
func TestJoinReader(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop()
// Create a table where each row is:
//
// | a | b | sum | s |
// |-----------------------------------------------------------------|
// | rowId/10 | rowId%10 | rowId/10 + rowId%10 | IntToEnglish(rowId) |
aFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row / 10))
}
bFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row % 10))
}
sumFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row/10 + row%10))
}
sqlutils.CreateTable(t, sqlDB, "t",
"a INT, b INT, sum INT, s STRING, PRIMARY KEY (a,b), INDEX bs (b,s)",
99,
sqlutils.ToRowFn(aFn, bFn, sumFn, sqlutils.RowEnglishFn))
td := sqlbase.GetTableDescriptor(kvDB, "test", "t")
testCases := []struct {
spec JoinReaderSpec
input [][]parser.Datum
expected string
}{
{
spec: JoinReaderSpec{
OutputColumns: []uint32{0, 1, 2},
},
input: [][]parser.Datum{
{aFn(2), bFn(2)},
{aFn(5), bFn(5)},
{aFn(10), bFn(10)},
{aFn(15), bFn(15)},
},
expected: "[[0 2 2] [0 5 5] [1 0 1] [1 5 6]]",
},
{
spec: JoinReaderSpec{
Filter: Expression{Expr: "$2 <= 5"}, // sum <= 5
OutputColumns: []uint32{3},
},
input: [][]parser.Datum{
{aFn(1), bFn(1)},
{aFn(25), bFn(25)},
{aFn(5), bFn(5)},
{aFn(21), bFn(21)},
{aFn(34), bFn(34)},
{aFn(13), bFn(13)},
{aFn(51), bFn(51)},
{aFn(50), bFn(50)},
},
expected: "[['one'] ['five'] ['two-one'] ['one-three'] ['five-zero']]",
},
}
for _, c := range testCases {
js := c.spec
js.Table = *td
txn := client.NewTxn(context.Background(), *kvDB)
flowCtx := FlowCtx{
Context: context.Background(),
evalCtx: &parser.EvalContext{},
txn: txn,
}
in := &RowBuffer{}
for _, row := range c.input {
encRow := make(sqlbase.EncDatumRow, len(row))
for i, d := range row {
encRow[i].SetDatum(sqlbase.ColumnType_INT, d)
}
in.rows = append(in.rows, encRow)
}
out := &RowBuffer{}
jr, err := newJoinReader(&flowCtx, &js, in, out)
if err != nil {
t.Fatal(err)
}
jr.Run(nil)
if out.err != nil {
t.Fatal(out.err)
}
if !in.done {
t.Fatal("joinReader stopped accepting rows")
}
if !out.closed {
t.Fatalf("output RowReceiver not closed")
}
if result := out.rows.String(); result != c.expected {
t.Errorf("invalid results: %s, expected %s'", result, c.expected)
}
}
}
// Prepare returns the result types of the given statement. pinfo may
// contain partial type information for placeholders. Prepare will
// populate the missing types. The column result types are returned (or
// nil if there are no results).
func (e *Executor) Prepare(
query string, session *Session, pinfo parser.PlaceholderTypes,
) (ResultColumns, error) {
log.VEventf(session.Ctx(), 2, "preparing: %s", query)
var p parser.Parser
stmts, err := p.Parse(query, parser.Syntax(session.Syntax))
if err != nil {
return nil, err
}
switch len(stmts) {
case 0:
return nil, nil
case 1:
// ignore
default:
return nil, errors.Errorf("expected 1 statement, but found %d", len(stmts))
}
stmt := stmts[0]
if err = pinfo.ProcessPlaceholderAnnotations(stmt); err != nil {
return nil, err
}
protoTS, err := isAsOf(&session.planner, stmt, e.cfg.Clock.Now())
if err != nil {
return nil, err
}
session.planner.resetForBatch(e)
session.planner.semaCtx.Placeholders.SetTypes(pinfo)
session.planner.evalCtx.PrepareOnly = true
// Prepare needs a transaction because it needs to retrieve db/table
// descriptors for type checking.
// TODO(andrei): is this OK? If we're preparing as part of a SQL txn, how do
// we check that they're reading descriptors consistent with the txn in which
// they'll be used?
txn := client.NewTxn(session.Ctx(), *e.cfg.DB)
txn.Proto.Isolation = session.DefaultIsolationLevel
session.planner.setTxn(txn)
defer session.planner.setTxn(nil)
if protoTS != nil {
session.planner.avoidCachedDescriptors = true
defer func() {
session.planner.avoidCachedDescriptors = false
}()
setTxnTimestamps(txn, *protoTS)
}
plan, err := session.planner.prepare(stmt)
if err != nil {
return nil, err
}
if plan == nil {
return nil, nil
}
defer plan.Close()
cols := plan.Columns()
for _, c := range cols {
if err := checkResultType(c.Typ); err != nil {
return nil, err
}
}
return cols, nil
}
func (ds *ServerImpl) setupTxn(ctx context.Context, txnProto *roachpb.Transaction) *client.Txn {
txn := client.NewTxn(ctx, *ds.DB)
// TODO(radu): we should sanity check some of these fields
txn.Proto = *txnProto
return txn
}
func (flowCtx *FlowCtx) setupTxn(ctx context.Context) *client.Txn {
txn := client.NewTxn(ctx, *flowCtx.clientDB)
txn.Proto = *flowCtx.txnProto
return txn
}
func TestServer(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop()
conn, err := s.RPCContext().GRPCDial(s.ServingAddr())
if err != nil {
t.Fatal(err)
}
r := sqlutils.MakeSQLRunner(t, sqlDB)
r.Exec(`CREATE DATABASE test`)
r.Exec(`CREATE TABLE test.t (a INT PRIMARY KEY, b INT)`)
r.Exec(`INSERT INTO test.t VALUES (1, 10), (2, 20), (3, 30)`)
td := sqlbase.GetTableDescriptor(kvDB, "test", "t")
ts := TableReaderSpec{
Table: *td,
IndexIdx: 0,
Reverse: false,
Spans: nil,
Filter: Expression{Expr: "$0 != 2"}, // a != 2
OutputColumns: []uint32{0, 1}, // a
}
txn := client.NewTxn(context.Background(), *kvDB)
req := &SetupFlowRequest{Txn: txn.Proto}
req.Flow = FlowSpec{
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &ts},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{{Mailbox: &MailboxSpec{SimpleResponse: true}}},
}},
}},
}
distSQLClient := NewDistSQLClient(conn)
stream, err := distSQLClient.RunSimpleFlow(context.Background(), req)
if err != nil {
t.Fatal(err)
}
var decoder StreamDecoder
var rows sqlbase.EncDatumRows
for {
msg, err := stream.Recv()
if err != nil {
if err == io.EOF {
break
}
t.Fatal(err)
}
err = decoder.AddMessage(msg)
if err != nil {
t.Fatal(err)
}
rows = testGetDecodedRows(t, &decoder, rows)
}
if done, trailerErr := decoder.IsDone(); !done {
t.Fatal("stream not done")
} else if trailerErr != nil {
t.Fatal("error in the stream trailer:", trailerErr)
}
str := rows.String()
expected := "[[1 10] [3 30]]"
if str != expected {
t.Errorf("invalid results: %s, expected %s'", str, expected)
}
}
