Skip to content

Commit 43e0360

Browse files
austinvallerainkwanstephybun
authored
statestore: Implement Write for state stores (#1270)
* initial implementation pair session Co-authored-by: Rain Kwan <[email protected]> * satisfy interface * add tests + update logic * package docs * update copywrite headers * Update statestore/write.go Co-authored-by: stephybun <[email protected]> * update error messaging --------- Co-authored-by: Rain Kwan <[email protected]> Co-authored-by: stephybun <[email protected]>
1 parent e642d51 commit 43e0360

File tree

12 files changed

+1256
-10
lines changed

12 files changed

+1256
-10
lines changed

‎internal/fwserver/server_writestatebytes.go‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
// Copyright IBM Corp. 2021, 2026
2+
// SPDX-License-Identifier: MPL-2.0
3+
4+
package fwserver
5+
6+
import (
7+
"context"
8+
9+
"github.com/hashicorp/terraform-plugin-framework/diag"
10+
"github.com/hashicorp/terraform-plugin-framework/internal/logging"
11+
"github.com/hashicorp/terraform-plugin-framework/statestore"
12+
)
13+
14+
type WriteStateBytesRequest struct {
15+
StateID string
16+
StateBytes []byte
17+
StateStore statestore.StateStore
18+
}
19+
20+
type WriteStateBytesResponse struct {
21+
Diagnostics diag.Diagnostics
22+
}
23+
24+
// WriteStateBytes implements the framework server WriteStateBytes RPC.
25+
func (s *Server) WriteStateBytes(ctx context.Context, req *WriteStateBytesRequest, resp *WriteStateBytesResponse) {
26+
if req == nil {
27+
return
28+
}
29+
30+
if stateStoreWithConfigure, ok := req.StateStore.(statestore.StateStoreWithConfigure); ok {
31+
logging.FrameworkTrace(ctx, "StateStore implements StateStoreWithConfigure")
32+
33+
configureReq := statestore.ConfigureRequest{
34+
StateStoreData: s.StateStoreConfigureData.StateStoreConfigureData,
35+
}
36+
configureResp := statestore.ConfigureResponse{}
37+
38+
logging.FrameworkTrace(ctx, "Calling provider defined StateStore Configure")
39+
stateStoreWithConfigure.Configure(ctx, configureReq, &configureResp)
40+
logging.FrameworkTrace(ctx, "Called provider defined StateStore Configure")
41+
42+
resp.Diagnostics.Append(configureResp.Diagnostics...)
43+
44+
if resp.Diagnostics.HasError() {
45+
return
46+
}
47+
}
48+
49+
writeReq := statestore.WriteRequest{
50+
StateID: req.StateID,
51+
StateBytes: req.StateBytes,
52+
}
53+
54+
writeResp := statestore.WriteResponse{}
55+
56+
logging.FrameworkTrace(ctx, "Calling provider defined StateStore Write")
57+
req.StateStore.Write(ctx, writeReq, &writeResp)
58+
logging.FrameworkTrace(ctx, "Called provider defined StateStore Write")
59+
60+
resp.Diagnostics.Append(writeResp.Diagnostics...)
61+
}

‎internal/fwserver/server_writestatebytes_test.go‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
// Copyright IBM Corp. 2021, 2026
2+
// SPDX-License-Identifier: MPL-2.0
3+
4+
package fwserver_test
5+
6+
import (
7+
"context"
8+
"fmt"
9+
"testing"
10+
11+
"github.com/google/go-cmp/cmp"
12+
"github.com/hashicorp/terraform-plugin-framework/diag"
13+
"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
14+
"github.com/hashicorp/terraform-plugin-framework/internal/testing/testprovider"
15+
"github.com/hashicorp/terraform-plugin-framework/statestore"
16+
)
17+
18+
func TestServerWriteStateBytes(t *testing.T) {
19+
t.Parallel()
20+
21+
testCases := map[string]struct {
22+
server *fwserver.Server
23+
request *fwserver.WriteStateBytesRequest
24+
expectedResponse *fwserver.WriteStateBytesResponse
25+
}{
26+
"nil": {
27+
server: &fwserver.Server{
28+
Provider: &testprovider.Provider{},
29+
},
30+
expectedResponse: &fwserver.WriteStateBytesResponse{},
31+
},
32+
"request": {
33+
server: &fwserver.Server{
34+
Provider: &testprovider.Provider{},
35+
},
36+
request: &fwserver.WriteStateBytesRequest{
37+
StateID: "test-state-123",
38+
StateBytes: []byte(`{"version": 4, "terraform_version": "1.15.0"}`),
39+
StateStore: &testprovider.StateStore{
40+
WriteMethod: func(ctx context.Context, req statestore.WriteRequest, resp *statestore.WriteResponse) {
41+
if req.StateID != "test-state-123" {
42+
resp.Diagnostics.AddError(
43+
"Unexpected req.StateID",
44+
fmt.Sprintf("expected \"test-state-123\", got: %q", req.StateID),
45+
)
46+
return
47+
}
48+
49+
if string(req.StateBytes) != `{"version": 4, "terraform_version": "1.15.0"}` {
50+
resp.Diagnostics.AddError(
51+
"Unexpected req.StateBytes",
52+
fmt.Sprintf("expected \"test-state-123\", got: %q", string(req.StateBytes)),
53+
)
54+
return
55+
}
56+
},
57+
},
58+
},
59+
expectedResponse: &fwserver.WriteStateBytesResponse{},
60+
},
61+
"statestore-configure-data": {
62+
server: &fwserver.Server{
63+
StateStoreConfigureData: fwserver.StateStoreConfigureData{
64+
StateStoreConfigureData: "test-statestore-configure-value",
65+
},
66+
Provider: &testprovider.Provider{},
67+
},
68+
request: &fwserver.WriteStateBytesRequest{
69+
StateStore: &testprovider.StateStoreWithConfigure{
70+
StateStore: &testprovider.StateStore{},
71+
ConfigureMethod: func(ctx context.Context, req statestore.ConfigureRequest, resp *statestore.ConfigureResponse) {
72+
stateStoreData, ok := req.StateStoreData.(string)
73+
74+
if !ok {
75+
resp.Diagnostics.AddError(
76+
"Unexpected ConfigureRequest.StateStoreData",
77+
fmt.Sprintf("Expected string, got: %T", req.StateStoreData),
78+
)
79+
return
80+
}
81+
82+
if stateStoreData != "test-statestore-configure-value" {
83+
resp.Diagnostics.AddError(
84+
"Unexpected ConfigureRequest.StateStoreData",
85+
fmt.Sprintf("Expected test-statestore-configure-value, got: %q", stateStoreData),
86+
)
87+
}
88+
},
89+
},
90+
},
91+
expectedResponse: &fwserver.WriteStateBytesResponse{},
92+
},
93+
"response-diagnostics": {
94+
server: &fwserver.Server{
95+
Provider: &testprovider.Provider{},
96+
},
97+
request: &fwserver.WriteStateBytesRequest{
98+
StateStore: &testprovider.StateStore{
99+
WriteMethod: func(ctx context.Context, req statestore.WriteRequest, resp *statestore.WriteResponse) {
100+
resp.Diagnostics.AddWarning("warning summary", "warning detail")
101+
resp.Diagnostics.AddError("error summary", "error detail")
102+
},
103+
},
104+
},
105+
expectedResponse: &fwserver.WriteStateBytesResponse{
106+
Diagnostics: diag.Diagnostics{
107+
diag.NewWarningDiagnostic(
108+
"warning summary",
109+
"warning detail",
110+
),
111+
diag.NewErrorDiagnostic(
112+
"error summary",
113+
"error detail",
114+
),
115+
},
116+
},
117+
},
118+
}
119+
120+
for name, testCase := range testCases {
121+
t.Run(name, func(t *testing.T) {
122+
t.Parallel()
123+
124+
response := &fwserver.WriteStateBytesResponse{}
125+
testCase.server.WriteStateBytes(context.Background(), testCase.request, response)
126+
127+
if diff := cmp.Diff(response, testCase.expectedResponse); diff != "" {
128+
t.Errorf("unexpected difference: %s", diff)
129+
}
130+
})
131+
}
132+
}

‎internal/proto6server/serve.go‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,3 @@ func (s *Server) DeleteState(ctx context.Context, req *tfprotov6.DeleteStateRequ
6161
// TODO: Implement in a separate PR for just workspace RPCs
6262
panic("DeleteState not implemented")
6363
}
64-
65-
func (s *Server) WriteStateBytes(ctx context.Context, req *tfprotov6.WriteStateBytesStream) (*tfprotov6.WriteStateBytesResponse, error) {
66-
// TODO: Implement in a separate PR for just WriteStateBytes
67-
panic("WriteStateBytes not implemented")
68-
}

‎internal/proto6server/server_readstatebytes.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,17 +53,15 @@ func (s *Server) ReadStateBytes(ctx context.Context, proto6Req *tfprotov6.ReadSt
5353
}
5454

5555
// If ConfigureStateStore isn't called prior to ReadStateBytes
56-
if int(s.FrameworkServer.StateStoreConfigureData.ServerCapabilities.ChunkSize) == 0 {
56+
if int(s.FrameworkServer.StateStoreConfigureData.ServerCapabilities.ChunkSize) <= 0 {
5757
return &tfprotov6.ReadStateBytesStream{
5858
Chunks: func(push func(chunk tfprotov6.ReadStateByteChunk) bool) {
5959
push(tfprotov6.ReadStateByteChunk{
6060
Diagnostics: []*tfprotov6.Diagnostic{
6161
{
6262
Severity: tfprotov6.DiagnosticSeverityError,
6363
Summary: "Error reading state",
64-
Detail: fmt.Sprintf("No chunk size received from Terraform while reading state data for %s. This is a bug and should be reported.",
65-
proto6Req.StateID,
66-
),
64+
Detail: "The provider server does not have a chunk size configured. This is a bug in either Terraform core or terraform-plugin-framework and should be reported.",
6765
},
6866
},
6967
})

‎internal/proto6server/server_readstatebytes_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ func TestServerReadStateBytes(t *testing.T) {
203203
{
204204
Severity: tfprotov6.DiagnosticSeverityError,
205205
Summary: "Error reading state",
206-
Detail: "No chunk size received from Terraform while reading state data for test_statestore. This is a bug and should be reported."},
206+
Detail: "The provider server does not have a chunk size configured. This is a bug in either Terraform core or terraform-plugin-framework and should be reported."},
207207
},
208208
},
209209
},

‎internal/proto6server/server_writestatebytes.go‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
// Copyright IBM Corp. 2021, 2026
2+
// SPDX-License-Identifier: MPL-2.0
3+
4+
package proto6server
5+
6+
import (
7+
"bytes"
8+
"context"
9+
"fmt"
10+
11+
"github.com/hashicorp/terraform-plugin-framework/internal/fwserver"
12+
"github.com/hashicorp/terraform-plugin-framework/internal/logging"
13+
"github.com/hashicorp/terraform-plugin-framework/internal/toproto6"
14+
"github.com/hashicorp/terraform-plugin-go/tfprotov6"
15+
"github.com/hashicorp/terraform-plugin-log/tflog"
16+
"github.com/hashicorp/terraform-plugin-log/tfsdklog"
17+
)
18+
19+
func (s *Server) WriteStateBytes(ctx context.Context, proto6Req *tfprotov6.WriteStateBytesStream) (*tfprotov6.WriteStateBytesResponse, error) {
20+
ctx = s.registerContext(ctx)
21+
ctx = logging.InitContext(ctx)
22+
23+
fwResp := &fwserver.WriteStateBytesResponse{}
24+
25+
var typeName string
26+
var stateID string
27+
var stateBuffer bytes.Buffer
28+
29+
// Chunk size is set on the server during the ConfigureStateStore RPC
30+
configuredChunkSize := s.FrameworkServer.StateStoreConfigureData.ServerCapabilities.ChunkSize
31+
if configuredChunkSize <= 0 {
32+
fwResp.Diagnostics.AddError(
33+
"Error Writing State",
34+
"The provider server does not have a chunk size configured. This is a bug in either Terraform core or terraform-plugin-framework and should be reported.",
35+
)
36+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
37+
}
38+
39+
// This field will be collected from one of the chunks
40+
var expectedTotalLength int64 = 0
41+
42+
for chunk, diags := range proto6Req.Chunks {
43+
// Any diagnostics here are either a GRPC communication error or invalid data from the client.
44+
if len(diags) > 0 {
45+
return &tfprotov6.WriteStateBytesResponse{Diagnostics: diags}, nil
46+
}
47+
48+
// Only the first chunk will have meta set
49+
if chunk.Meta != nil {
50+
typeName = chunk.Meta.TypeName
51+
stateID = chunk.Meta.StateID
52+
}
53+
54+
if chunk.Range.End < chunk.TotalLength-1 {
55+
// Ensure each chunk (except the last) exactly match the configured size.
56+
if int64(len(chunk.Bytes)) != configuredChunkSize {
57+
fwResp.Diagnostics.AddError(
58+
"Error Writing State",
59+
fmt.Sprintf("Unexpected chunk of size %d was received from Terraform, expected chunk size was %d. This is a bug in Terraform core and should be reported.", len(chunk.Bytes), configuredChunkSize),
60+
)
61+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
62+
}
63+
} else {
64+
// Ensure the last chunk is within the configured size.
65+
if int64(len(chunk.Bytes)) > configuredChunkSize {
66+
fwResp.Diagnostics.AddError(
67+
"Error Writing State",
68+
fmt.Sprintf("Unexpected final chunk of size %d was received from Terraform, which exceeds the configured chunk size of %d. This is a bug in Terraform core and should be reported.", len(chunk.Bytes), configuredChunkSize),
69+
)
70+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
71+
}
72+
}
73+
74+
if expectedTotalLength == 0 {
75+
expectedTotalLength = chunk.TotalLength
76+
}
77+
78+
stateBuffer.Write(chunk.Bytes)
79+
}
80+
81+
// MAINTAINER NOTE: Typically these fields are set in terraform-plugin-go (see link below), however because
82+
// the type name is not extracted until the stream is consumed, it's easier to set the logger fields here.
83+
//
84+
// https://github.com/hashicorp/terraform-plugin-go/blob/14fe65ea923b5e306dbb4f67f2bb861f74b9e3ec/internal/logging/context.go#L112-L119
85+
ctx = tfsdklog.SetField(ctx, "tf_state_store_type", typeName)
86+
ctx = tfsdklog.SubsystemSetField(ctx, "proto", "tf_state_store_type", typeName)
87+
ctx = tflog.SetField(ctx, "tf_state_store_type", typeName)
88+
89+
if stateBuffer.Len() == 0 {
90+
fwResp.Diagnostics.AddError(
91+
"Error Writing State",
92+
"No state data was received from Terraform. This is a bug in Terraform core and should be reported.",
93+
)
94+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
95+
}
96+
97+
if int64(stateBuffer.Len()) != expectedTotalLength {
98+
fwResp.Diagnostics.AddError(
99+
"Error Writing State",
100+
fmt.Sprintf("Unexpected size of state data received from Terraform, got: %d, expected: %d. This is a bug in Terraform core and should be reported.", stateBuffer.Len(), expectedTotalLength),
101+
)
102+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
103+
}
104+
105+
if stateID == "" {
106+
fwResp.Diagnostics.AddError(
107+
"Error Writing State",
108+
"No state ID was received from Terraform. This is a bug in Terraform core and should be reported.",
109+
)
110+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
111+
}
112+
113+
stateStore, diags := s.FrameworkServer.StateStore(ctx, typeName)
114+
115+
fwResp.Diagnostics.Append(diags...)
116+
117+
if fwResp.Diagnostics.HasError() {
118+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
119+
}
120+
121+
fwReq := &fwserver.WriteStateBytesRequest{
122+
StateStore: stateStore,
123+
StateID: stateID,
124+
StateBytes: stateBuffer.Bytes(),
125+
}
126+
127+
s.FrameworkServer.WriteStateBytes(ctx, fwReq, fwResp)
128+
129+
return toproto6.WriteStateBytesResponse(ctx, fwResp), nil
130+
}

0 commit comments

Comments
 (0)