code.vegaprotocol.io/vega@v0.79.0/datanode/api/export.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package api
    17  
    18  import (
    19  	"bytes"
    20  	"fmt"
    21  
    22  	"google.golang.org/genproto/googleapis/api/httpbody"
    23  )
    24  
    25  type httpBodyWriter struct {
    26  	chunkSize   int
    27  	contentType string
    28  	buf         *bytes.Buffer
    29  	stream      tradingDataServiceExportServer
    30  }
    31  
    32  type tradingDataServiceExportServer interface {
    33  	Send(*httpbody.HttpBody) error
    34  }
    35  
    36  func (w *httpBodyWriter) Write(p []byte) (n int, err error) {
    37  	n = len(p)
    38  	w.buf.Write(p)
    39  
    40  	if w.buf.Len() >= w.chunkSize {
    41  		if err := w.sendChunk(); err != nil {
    42  			return 0, err
    43  		}
    44  	}
    45  
    46  	return n, nil
    47  }
    48  
    49  func (w *httpBodyWriter) sendChunk() error {
    50  	msg := &httpbody.HttpBody{
    51  		ContentType: w.contentType,
    52  		Data:        w.buf.Bytes(),
    53  	}
    54  
    55  	if err := w.stream.Send(msg); err != nil {
    56  		return fmt.Errorf("error sending chunk: %w", err)
    57  	}
    58  
    59  	w.buf.Reset()
    60  	return nil
    61  }
    62  
    63  func (w *httpBodyWriter) Close() error {
    64  	if w.buf.Len() > 0 {
    65  		return w.sendChunk()
    66  	}
    67  	return nil
    68  }