github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/workloadccl/format/sstable.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package format
    10  
    11  import (
    12  	"context"
    13  	"fmt"
    14  	"time"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/ccl/importccl"
    17  	"github.com/cockroachdb/cockroach/pkg/keys"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/row"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    23  	"github.com/cockroachdb/cockroach/pkg/storage"
    24  	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
    25  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    26  	"github.com/cockroachdb/cockroach/pkg/workload"
    27  	"github.com/cockroachdb/errors"
    28  )
    29  
    30  // ToTableDescriptor returns the corresponding TableDescriptor for a workload
    31  // Table.
    32  func ToTableDescriptor(
    33  	t workload.Table, tableID sqlbase.ID, ts time.Time,
    34  ) (*sqlbase.TableDescriptor, error) {
    35  	ctx := context.Background()
    36  	stmt, err := parser.ParseOne(fmt.Sprintf(`CREATE TABLE "%s" %s`, t.Name, t.Schema))
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	createTable, ok := stmt.AST.(*tree.CreateTable)
    41  	if !ok {
    42  		return nil, errors.Errorf("expected *tree.CreateTable got %T", stmt)
    43  	}
    44  	const parentID sqlbase.ID = keys.MaxReservedDescID
    45  	tableDesc, err := importccl.MakeSimpleTableDescriptor(
    46  		ctx, nil /* settings */, createTable, parentID, tableID, importccl.NoFKs, ts.UnixNano())
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  	return &tableDesc.TableDescriptor, nil
    51  }
    52  
    53  // ToSSTable constructs a single sstable with the kvs necessary to represent a
    54  // workload.Table as a CockroachDB SQL table. This sstable is suitable for
    55  // handing to AddSSTable or RocksDB's IngestExternalFile.
    56  //
    57  // TODO(dan): Finally remove sampledataccl in favor of this.
    58  func ToSSTable(t workload.Table, tableID sqlbase.ID, ts time.Time) ([]byte, error) {
    59  	ctx := context.Background()
    60  	tableDesc, err := ToTableDescriptor(t, tableID, ts)
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	kvCh := make(chan row.KVBatch)
    66  	wc := importccl.NewWorkloadKVConverter(
    67  		0, tableDesc, t.InitialRows, 0, t.InitialRows.NumBatches, kvCh)
    68  
    69  	g := ctxgroup.WithContext(ctx)
    70  	g.GoCtx(func(ctx context.Context) error {
    71  		defer close(kvCh)
    72  		evalCtx := &tree.EvalContext{SessionData: &sessiondata.SessionData{}}
    73  		return wc.Worker(ctx, evalCtx)
    74  	})
    75  	var sst []byte
    76  	g.GoCtx(func(ctx context.Context) error {
    77  		sstTS := hlc.Timestamp{WallTime: ts.UnixNano()}
    78  		sstFile := &storage.MemFile{}
    79  		sw := storage.MakeIngestionSSTWriter(sstFile)
    80  		defer sw.Close()
    81  		for kvBatch := range kvCh {
    82  			for _, kv := range kvBatch.KVs {
    83  				mvccKey := storage.MVCCKey{Timestamp: sstTS, Key: kv.Key}
    84  				if err := sw.Put(mvccKey, kv.Value.RawBytes); err != nil {
    85  					return err
    86  				}
    87  			}
    88  		}
    89  		err = sw.Finish()
    90  		sst = sstFile.Data()
    91  		return err
    92  	})
    93  	if err := g.Wait(); err != nil {
    94  		return nil, err
    95  	}
    96  
    97  	return sst, nil
    98  }