github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/table/io.go (about) 1 // Copyright 2019 Dolthub, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package table 16 17 import ( 18 "context" 19 "errors" 20 "io" 21 22 "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" 23 "github.com/dolthub/dolt/go/libraries/doltcore/row" 24 "github.com/dolthub/dolt/go/libraries/doltcore/schema" 25 "github.com/dolthub/dolt/go/store/types" 26 ) 27 28 // GetRow returns a row from |tbl| corresponding to |key| if it exists. 29 func GetRow(ctx context.Context, tbl *doltdb.Table, sch schema.Schema, key types.Tuple) (r row.Row, ok bool, err error) { 30 rowMap, err := tbl.GetRowData(ctx) 31 if err != nil { 32 return nil, false, err 33 } 34 35 var fields types.Value 36 fields, ok, err = rowMap.MaybeGet(ctx, key) 37 if err != nil || !ok { 38 return nil, ok, err 39 } 40 41 r, err = row.FromNoms(sch, key, fields.(types.Tuple)) 42 return 43 } 44 45 // PipeRows will read a row from given TableReader and write it to the provided TableWriter. It will do this 46 // for every row until the TableReader's ReadRow method returns io.EOF or encounters an error in either reading 47 // or writing. The caller will need to handle closing the tables as necessary. If contOnBadRow is true, errors reading 48 // or writing will be ignored and the pipe operation will continue. 49 // 50 // Returns a tuple: (number of rows written, number of errors ignored, error). In the case that err is non-nil, the 51 // row counter fields in the tuple will be set to -1. 52 func PipeRows(ctx context.Context, rd TableReader, wr TableWriter, contOnBadRow bool) (int, int, error) { 53 var numBad, numGood int 54 for { 55 r, err := rd.ReadRow(ctx) 56 57 if err != nil && err != io.EOF { 58 if IsBadRow(err) && contOnBadRow { 59 numBad++ 60 continue 61 } 62 63 return -1, -1, err 64 } else if err == io.EOF && r == nil { 65 break 66 } else if r == nil { 67 // row equal to nil should 68 return -1, -1, errors.New("reader returned nil row with err==nil") 69 } 70 71 err = wr.WriteRow(ctx, r) 72 73 if err != nil { 74 return -1, -1, err 75 } else { 76 numGood++ 77 } 78 } 79 80 return numGood, numBad, nil 81 } 82 83 // ReadAllRows reads all rows from a TableReader and returns a slice containing those rows. Usually this is used 84 // for testing, or with very small data sets. 85 func ReadAllRows(ctx context.Context, rd TableReader, contOnBadRow bool) ([]row.Row, int, error) { 86 var rows []row.Row 87 var err error 88 89 badRowCount := 0 90 for { 91 var r row.Row 92 r, err = rd.ReadRow(ctx) 93 94 if err != nil && err != io.EOF || r == nil { 95 if IsBadRow(err) { 96 badRowCount++ 97 98 if contOnBadRow { 99 continue 100 } 101 } 102 103 break 104 } 105 106 rows = append(rows, r) 107 } 108 109 if err == nil || err == io.EOF { 110 return rows, badRowCount, nil 111 } 112 113 return nil, badRowCount, err 114 }