PSS: Implement ReadStateBytes + WriteStateBytes (#37440)

* Implement ReadStateBytes + WriteStateBytes

* [WIP] -  Testing ReadStateBytes and WriteStateBytes (#37464)

* Fix nil pointer error

* Add WIP test for ReadStateBytes

* Move test mock to separate testing file

* Update mock to send unexpected EOF when there's a problem returning data and it's not a true EOF

* Add test case for when length != expected length

* Add test for when trying to read state from a store type that doesn't exist

* Change symbol names to lowercase

* Add ability to force a diagnostic to be returned from `mockReadStateBytesClient`'s `Recv` method

* Add test showing error diagnostics raised by the ReadStateBytes client are returned

* Add missing header

* Simplify mock by using an embedded type

* Rename `mockOpts` to `mockReadStateBytesOpts`

* Update existing tests to assert what arguments are passed to the RPC method call

* Add mock WriteStateBytesClient which uses `go.uber.org/mock/gomock` to enable assertions about calls to Send

* Add a test for WriteStateBytes that makes assertions about calls to the Send method

* Update test case to explicitly test writing data smaller than the chunk size

* Implement chunking in WriteStateBytes, add test case to assert expected chunking behaviour

* Add generated mock for Provider_WriteStateBytesClient in protocol v6

* Update tests to use new `MockProvider_WriteStateBytesClient`, remove handwritten mock

* Update code comments in test

* Add tests for diagnostics and errors returned during WriteStateBytes

* Add generated mock for Provider_ReadStateBytesClient in protocol v6, replace old mock

* Add test case for grpc errors in ReadStateBytes, fix how error is returned

* Typo in comment

* Add missing warning test, rename some test cases

* Update proto file definition of Read/WriteStateBytes RPCs (#37529)

* Update Read/WriteStateBytes RPCs to match hashicorp/terraform-plugin-go#531

* Run `make protobuf`

* Run `make generate`

* Update use of `proto.ReadStateBytes_ResponseChunk` in tests

* Fix how diagnostics are handled alongside EOF error, update ReadStateBytes test

* More fixes - test setup was incorrect

I think? I assume that a response would be returned full of zero-values when EOF is encountered.

* WIP - avoid crash if chunk is nil

* Sarah's updates to radek/pss-read-write (#37642)

* Update code to not expect a chunk when EOF encountered

* Return early if any grpc errors are encountered during ReadStateBytes

* Close the stream with CloseSend once everything's read without error. Add test case about handling grpc errors from CloseSend.

* Fix test case about warnings: We would expect to receive a chunk with data alongside the warning and have a normal closing of the stream after EOF

* Add log line, remove unneeded type info

* Implement configurable state chunk size

* handle metadata in WriteStateBytes correctly

* validate chunk sizes received from provider

* ReadStateBytes: avoid early return on warnings

---------

Co-authored-by: Sarah French <[email protected]>

Author: radeksimko

docs/plugin-protocol/tfplugin6.proto

@@ -390,6 +390,11 @@ service Provider {
     // ConfigureStateStore configures the state store, such as S3 connection in the context of already configured provider
     rpc ConfigureStateStore(ConfigureStateStore.Request) returns (ConfigureStateStore.Response);
 
+    // ReadStateBytes streams byte chunks of a given state file from a state store
+    rpc ReadStateBytes(ReadStateBytes.Request) returns (stream ReadStateBytes.Response);
+    // WriteStateBytes streams byte chunks of a given state file into a state store
+    rpc WriteStateBytes(stream WriteStateBytes.RequestChunk) returns (WriteStateBytes.Response);
+
     // GetStates returns a list of all states (i.e. CE workspaces) managed by a given state store
     rpc GetStates(GetStates.Request) returns (GetStates.Response);
     // DeleteState instructs a given state store to delete a specific state (i.e. a CE workspace)
@@ -896,7 +901,6 @@ message ValidateListResourceConfig {
     }
 }
 
-
 message ValidateStateStore {
     message Request {
         string type_name = 1;
@@ -911,12 +915,59 @@ message ConfigureStateStore {
     message Request {
         string type_name = 1;
         DynamicValue config = 2;
+        StateStoreClientCapabilities capabilities = 3;
     }
     message Response {
         repeated Diagnostic diagnostics = 1;
+        StateStoreServerCapabilities capabilities = 2;
     }
 }
 
+message StateStoreClientCapabilities {
+    int64 chunk_size = 1; // suggested chunk size by Core
+}
+
+message StateStoreServerCapabilities {
+    int64 chunk_size = 1; // chosen chunk size by plugin
+}
+
+message ReadStateBytes {
+    message Request {
+        string type_name = 1;
+        string state_id = 2;
+    }
+    message Response {
+        bytes bytes = 1;
+        int64 total_length = 2;
+        StateRange range = 3;
+        repeated Diagnostic diagnostics = 4;
+    }
+}
+
+message WriteStateBytes {
+    message RequestChunk {
+        // meta is sent with the first chunk only
+        optional RequestChunkMeta meta = 1;
+
+        bytes bytes = 2;
+        int64 total_length = 3;
+        StateRange range = 4;
+    }
+    message Response {
+        repeated Diagnostic diagnostics = 1;
+    }
+}
+
+message RequestChunkMeta {
+    string type_name = 1;
+    string state_id = 2;
+}
+
+message StateRange {
+    int64 start = 1;
+    int64 end = 2;
+}
+
 message GetStates {
 	message Request {
 		string type_name = 1;

internal/builtin/providers/terraform/provider.go

@@ -295,6 +295,18 @@ func (p *Provider) ConfigureStateStore(req providers.ConfigureStateStoreRequest)
 	return resp
 }
 
+func (p *Provider) ReadStateBytes(req providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
+	var resp providers.ReadStateBytesResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
+}
+
+func (p *Provider) WriteStateBytes(req providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
+	var resp providers.WriteStateBytesResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
+}
+
 func (p *Provider) GetStates(req providers.GetStatesRequest) providers.GetStatesResponse {
 	var resp providers.GetStatesResponse
 	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))

internal/grpcwrap/provider6.go

@@ -907,6 +907,14 @@ func (p *provider6) ConfigureStateStore(ctx context.Context, req *tfplugin6.Conf
 	panic("not implemented")
 }
 
+func (p *provider6) ReadStateBytes(req *tfplugin6.ReadStateBytes_Request, srv tfplugin6.Provider_ReadStateBytesServer) error {
+	panic("not implemented")
+}
+
+func (p *provider6) WriteStateBytes(srv tfplugin6.Provider_WriteStateBytesServer) error {
+	panic("not implemented")
+}
+
 func (p *provider6) GetStates(ctx context.Context, req *tfplugin6.GetStates_Request) (*tfplugin6.GetStates_Response, error) {
 	panic("not implemented")
 }

internal/plugin/grpc_provider.go

@@ -1444,6 +1444,14 @@ func (p *GRPCProvider) ConfigureStateStore(r providers.ConfigureStateStoreReques
 	panic("not implemented")
 }
 
+func (p *GRPCProvider) ReadStateBytes(r providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
+	panic("not implemented")
+}
+
+func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
+	panic("not implemented")
+}
+
 func (p *GRPCProvider) GetStates(r providers.GetStatesRequest) providers.GetStatesResponse {
 	panic("not implemented")
 }