implement WriteStateBytes

Author: radeksimko

tfprotov6/state_store.go

@@ -56,6 +56,16 @@ type ReadStateBytesStream struct {
 	Chunks iter.Seq[ReadStateByteChunk]
 }
 
+type WriteStateBytesStream struct {
+	Chunks iter.Seq[WriteStateByteChunk]
+}
+
+type WriteStateBytesResponse struct {
+	Diagnostics []*Diagnostic
+}
+
+type WriteStateByteChunk = StateByteChunk
+
 type ReadStateByteChunk struct {
 	StateByteChunk
 	Diagnostics []*Diagnostic

tfprotov6/tf6server/server.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"os/signal"
 	"regexp"
@@ -1645,6 +1646,78 @@ func (s *server) ReadStateBytes(protoReq *tfplugin6.ReadStateBytes_Request, prot
 	return nil
 }
 
+func (s *server) WriteStateBytes(srv grpc.ClientStreamingServer[tfplugin6.WriteStateBytes_RequestChunk, tfplugin6.WriteStateBytes_Response]) error {
+	rpc := "WriteStateBytes"
+	ctx := srv.Context()
+	ctx = s.loggingContext(ctx)
+	ctx = logging.RpcContext(ctx, rpc)
+	// ctx = logging.StateStoreContext(ctx, protoReq.TypeName)
+	ctx = s.stoppableContext(ctx)
+	// logging.ProtocolTrace(ctx, "Received request")
+	// defer logging.ProtocolTrace(ctx, "Served request")
+
+	ctx = tf6serverlogging.DownstreamRequest(ctx)
+
+	server, ok := s.downstream.(tfprotov6.StateStoreServer)
+	if !ok {
+		err := status.Error(codes.Unimplemented, "ProviderServer does not implement WriteStateBytes")
+		logging.ProtocolError(ctx, err.Error())
+		return err
+	}
+
+	iterator := func(yield func(tfprotov6.WriteStateByteChunk) bool) {
+		for {
+			chunk, err := srv.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				// attempt to send the error back to client
+				msgErr := srv.SendMsg(&tfplugin6.WriteStateBytes_Response{
+					Diagnostics: toproto.Diagnostics([]*tfprotov6.Diagnostic{
+						{
+							Severity: tfprotov6.DiagnosticSeverityError,
+							Summary:  "Writing state chunk failed",
+							Detail: fmt.Sprintf("Attempt to write a byte chunk of state %q to %q failed: %s",
+								chunk.StateId, chunk.TypeName, err),
+						},
+					}),
+				})
+				if msgErr != nil {
+					err := status.Error(codes.Unimplemented, "ProviderServer does not implement WriteStateBytes")
+					logging.ProtocolError(ctx, err.Error())
+					return
+				}
+				return
+			}
+
+			ok := yield(tfprotov6.WriteStateByteChunk{
+				Bytes:       chunk.Bytes,
+				TotalLength: chunk.TotalLength,
+				Range: tfprotov6.StateByteRange{
+					Start: chunk.Range.Start,
+					End:   chunk.Range.End,
+				},
+			})
+			if !ok {
+				return
+			}
+
+		}
+	}
+
+	resp, err := server.WriteStateBytes(ctx, &tfprotov6.WriteStateBytesStream{
+		Chunks: iterator,
+	})
+	if err != nil {
+		return err
+	}
+
+	return srv.SendAndClose(&tfplugin6.WriteStateBytes_Response{
+		Diagnostics: toproto.Diagnostics(resp.Diagnostics),
+	})
+}
+
 func (s *server) GetStates(ctx context.Context, protoReq *tfplugin6.GetStates_Request) (*tfplugin6.GetStates_Response, error) {
 	rpc := "GetStates"
 	ctx = s.loggingContext(ctx)