github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/util/writers/max_line_writer.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  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2016 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package writers
    23  
    24  import "io"
    25  
    26  var (
    27  	// MaxLinesErr is an instance of MaxLinesError that gets returned by
    28  	// Write() whenever the number of lines written has exceeded the number
    29  	// in |MaxLineWriter.MaxLines|.
    30  	MaxLinesErr = MaxLinesError{"Maximum number of lines written"}
    31  )
    32  
    33  // MaxLinesError is the type of error returned by Write() whenever the number
    34  // of lines written has exceeded the number in |MaxLineWriter.MaxLines|.
    35  type MaxLinesError struct {
    36  	msg string
    37  }
    38  
    39  func (e MaxLinesError) Error() string { return e.msg }
    40  
    41  // MaxLineWriter provides an io.Writer interface that counts the number of lines
    42  // that have been written. It will stop writing and returns an error if the
    43  // number of lines written exceeds the number specified in MaxLineWriter.NumLines.
    44  type MaxLineWriter struct {
    45  	Dest     io.Writer
    46  	MaxLines uint32
    47  	NumLines uint32
    48  }
    49  
    50  // Write() stops writing and returns an error if an attempt is made to write
    51  // any byte after |MaxLines| newLines have been written. For example, if MaxLines
    52  // is 1, all bytes will be written up to and including the 1st newline. If there
    53  // are any bytes in |data| after the 1st newline, an error will be returned.
    54  //
    55  // Callers can change the value of |w.MaxLines| before any call to Write().
    56  // Setting MaxLines to 0 will allow any number of newLines.
    57  func (w *MaxLineWriter) Write(data []byte) (int, error) {
    58  	if len(data) == 0 {
    59  		return 0, nil
    60  	}
    61  
    62  	checkMax := w.MaxLines > 0
    63  
    64  	if checkMax && w.NumLines >= w.MaxLines {
    65  		return 0, MaxLinesErr
    66  	}
    67  
    68  	var err error
    69  	byteCnt := 0
    70  
    71  	for i, b := range data {
    72  		if b == byte('\n') {
    73  			w.NumLines++
    74  			if checkMax && w.NumLines > w.MaxLines {
    75  				err = MaxLinesErr
    76  				break
    77  			}
    78  		} else if checkMax && w.NumLines >= w.MaxLines {
    79  			err = MaxLinesErr
    80  			break
    81  		}
    82  		byteCnt = i
    83  	}
    84  
    85  	cnt, err1 := w.Dest.Write(data[:byteCnt+1])
    86  	if err1 != nil {
    87  		return cnt, err1
    88  	}
    89  	return cnt, err
    90  }