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 }