implement ReadStateBytes

Author: radeksimko

tfprotov6/internal/fromproto/state_store.go

@@ -30,6 +30,17 @@ func ConfigureStateStoreRequest(in *tfplugin6.ConfigureStateStore_Request) *tfpr
 	}
 }
 
+func ReadStateBytesRequest(in *tfplugin6.ReadStateBytes_Request) *tfprotov6.ReadStateBytesRequest {
+	if in == nil {
+		return nil
+	}
+
+	return &tfprotov6.ReadStateBytesRequest{
+		TypeName: in.TypeName,
+		StateId:  in.StateId,
+	}
+}
+
 func GetStatesRequest(in *tfplugin6.GetStates_Request) *tfprotov6.GetStatesRequest {
 	if in == nil {
 		return nil

tfprotov6/internal/toproto/state_store.go

@@ -28,6 +28,22 @@ func ConfigureStateStore_Response(in *tfprotov6.ConfigureStateStoreResponse) *tf
 	}
 }
 
+func ReadStateBytes_Response(in *tfprotov6.ReadStateByteChunk) *tfplugin6.ReadStateBytes_Response {
+	if in == nil {
+		return nil
+	}
+
+	return &tfplugin6.ReadStateBytes_Response{
+		Diagnostics: Diagnostics(in.Diagnostics),
+		Bytes:       in.Bytes,
+		TotalLength: in.TotalLength,
+		Range: &tfplugin6.StateRange{
+			Start: in.Range.Start,
+			End:   in.Range.End,
+		},
+	}
+}
+
 func GetStates_Response(in *tfprotov6.GetStatesResponse) *tfplugin6.GetStates_Response {
 	if in == nil {
 		return nil

tfprotov6/state_store.go

@@ -3,7 +3,10 @@
 
 package tfprotov6
 
-import "context"
+import (
+	"context"
+	"iter"
+)
 
 // StateStoreServer is an interface containing the methods an list resource
 // implementation needs to fill.
@@ -14,6 +17,11 @@ type StateStoreServer interface {
 	// ConfigureStateStore configures the state store, such as S3 connection in the context of already configured provider
 	ConfigureStateStore(context.Context, *ConfigureStateStoreRequest) (*ConfigureStateStoreResponse, error)
 
+	// ReadStateBytes streams byte chunks of a given state file from a state store
+	ReadStateBytes(context.Context, *ReadStateBytesRequest) (*ReadStateBytesStream, error)
+
+	WriteStateBytes(context.Context, *WriteStateBytesStream) (*WriteStateBytesResponse, error)
+
 	// GetStates returns a list of all states (i.e. CE workspaces) managed by a given state store
 	GetStates(context.Context, *GetStatesRequest) (*GetStatesResponse, error)
 
@@ -39,6 +47,30 @@ type ConfigureStateStoreResponse struct {
 	Diagnostics []*Diagnostic
 }
 
+type ReadStateBytesRequest struct {
+	TypeName string
+	StateId  string
+}
+
+type ReadStateBytesStream struct {
+	Chunks iter.Seq[ReadStateByteChunk]
+}
+
+type ReadStateByteChunk struct {
+	StateByteChunk
+	Diagnostics []*Diagnostic
+}
+
+type StateByteChunk struct {
+	Bytes       []byte
+	TotalLength int64
+	Range       StateByteRange
+}
+
+type StateByteRange struct {
+	Start, End int64
+}
+
 type GetStatesRequest struct {
 	TypeName string
 }

tfprotov6/tf6server/server.go

@@ -1598,6 +1598,53 @@ func (s *server) ConfigureStateStore(ctx context.Context, protoReq *tfplugin6.Co
 	return protoResp, nil
 }
 
+func (s *server) ReadStateBytes(protoReq *tfplugin6.ReadStateBytes_Request, protoStream grpc.ServerStreamingServer[tfplugin6.ReadStateBytes_Response]) error {
+	rpc := "ReadStateBytes"
+	ctx := protoStream.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")
+
+	req := fromproto.ReadStateBytesRequest(protoReq)
+	logging.ProtocolData(ctx, s.protocolDataDir, rpc, "Request", "StateId", req.StateId)
+
+	ctx = tf6serverlogging.DownstreamRequest(ctx)
+
+	server, ok := s.downstream.(tfprotov6.StateStoreServer)
+	if !ok {
+		err := status.Error(codes.Unimplemented, "ProviderServer does not implement ReadStateBytes")
+		logging.ProtocolError(ctx, err.Error())
+		return err
+	}
+
+	stream, err := server.ReadStateBytes(ctx, req)
+	if err != nil {
+		logging.ProtocolError(ctx, "Error from downstream", map[string]interface{}{logging.KeyError: err})
+		return err
+	}
+
+	for chunk := range stream.Chunks {
+		select {
+		// TODO: check how interruptions are handled
+		case <-ctx.Done():
+			logging.ProtocolTrace(ctx, "all chunks sent")
+			return nil
+
+		default:
+			protoChunk := toproto.ReadStateBytes_Response(&chunk)
+			if err := protoStream.Send(protoChunk); err != nil {
+				logging.ProtocolError(ctx, "Error sending chunk", map[string]any{logging.KeyError: err})
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 func (s *server) GetStates(ctx context.Context, protoReq *tfplugin6.GetStates_Request) (*tfplugin6.GetStates_Response, error) {
 	rpc := "GetStates"
 	ctx = s.loggingContext(ctx)