Implement LockState & UnlockState provider methods (#37711)

Author: radeksimko

docs/plugin-protocol/tfplugin6.proto

@@ -395,6 +395,11 @@ service Provider {
     // WriteStateBytes streams byte chunks of a given state file into a state store
     rpc WriteStateBytes(stream WriteStateBytes.RequestChunk) returns (WriteStateBytes.Response);
 
+    // LockState locks a given state (i.e. CE workspace)
+    rpc LockState(LockState.Request) returns (LockState.Response);
+    // UnlockState unlocks a given state (i.e. CE workspace)
+    rpc UnlockState(UnlockState.Request) returns (UnlockState.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)
@@ -968,6 +973,30 @@ message StateRange {
     int64 end = 2;
 }
 
+message LockState {
+    message Request {
+        string type_name = 1;
+        string state_id = 2;
+        // operation represents an ongoing operation due to which lock is held (e.g. refresh, plan, apply)
+        string operation = 3;
+    }
+    message Response {
+        string lock_id = 1;
+        repeated Diagnostic diagnostics = 2;
+    }
+}
+
+message UnlockState {
+    message Request {
+        string type_name = 1;
+        string state_id = 2;
+        string lock_id = 3;
+    }
+    message Response {
+        repeated Diagnostic diagnostics = 1;
+    }
+}
+
 message GetStates {
 	message Request {
 		string type_name = 1;

internal/builtin/providers/terraform/provider.go

@@ -307,6 +307,18 @@ func (p *Provider) WriteStateBytes(req providers.WriteStateBytesRequest) provide
 	return resp
 }
 
+func (p *Provider) LockState(req providers.LockStateRequest) providers.LockStateResponse {
+	var resp providers.LockStateResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
+}
+
+func (p *Provider) UnlockState(req providers.UnlockStateRequest) providers.UnlockStateResponse {
+	var resp providers.UnlockStateResponse
+	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

@@ -915,6 +915,14 @@ func (p *provider6) WriteStateBytes(srv tfplugin6.Provider_WriteStateBytesServer
 	panic("not implemented")
 }
 
+func (p *provider6) LockState(ctx context.Context, req *tfplugin6.LockState_Request) (*tfplugin6.LockState_Response, error) {
+	panic("not implemented")
+}
+
+func (p *provider6) UnlockState(ctx context.Context, req *tfplugin6.UnlockState_Request) (*tfplugin6.UnlockState_Response, 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

@@ -1452,6 +1452,14 @@ func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) provi
 	panic("not implemented")
 }
 
+func (p *GRPCProvider) LockState(r providers.LockStateRequest) providers.LockStateResponse {
+	panic("not implemented")
+}
+
+func (p *GRPCProvider) UnlockState(r providers.UnlockStateRequest) providers.UnlockStateResponse {
+	panic("not implemented")
+}
+
 func (p *GRPCProvider) GetStates(r providers.GetStatesRequest) providers.GetStatesResponse {
 	panic("not implemented")
 }

internal/plugin6/grpc_provider.go

@@ -1738,6 +1738,67 @@ func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) (resp
 	return resp
 }
 
+func (p *GRPCProvider) LockState(r providers.LockStateRequest) (resp providers.LockStateResponse) {
+	logger.Trace("GRPCProvider.v6: LockState")
+
+	schema := p.GetProviderSchema()
+	if schema.Diagnostics.HasErrors() {
+		resp.Diagnostics = schema.Diagnostics
+		return resp
+	}
+
+	_, ok := schema.StateStores[r.TypeName]
+	if !ok {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown state store type %q", r.TypeName))
+		return resp
+	}
+
+	protoReq := &proto6.LockState_Request{
+		TypeName:  r.TypeName,
+		StateId:   r.StateId,
+		Operation: r.Operation,
+	}
+
+	protoResp, err := p.client.LockState(p.ctx, protoReq)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
+		return resp
+	}
+	resp.LockId = protoResp.LockId
+	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
+	return resp
+}
+
+func (p *GRPCProvider) UnlockState(r providers.UnlockStateRequest) (resp providers.UnlockStateResponse) {
+	logger.Trace("GRPCProvider.v6: UnlockState")
+
+	schema := p.GetProviderSchema()
+	if schema.Diagnostics.HasErrors() {
+		resp.Diagnostics = schema.Diagnostics
+		return resp
+	}
+
+	_, ok := schema.StateStores[r.TypeName]
+	if !ok {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("unknown state store type %q", r.TypeName))
+		return resp
+	}
+
+	protoReq := &proto6.UnlockState_Request{
+		TypeName: r.TypeName,
+		StateId:  r.StateId,
+		LockId:   r.LockId,
+	}
+
+	protoResp, err := p.client.UnlockState(p.ctx, protoReq)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
+		return resp
+	}
+	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
+	return resp
+}
+
 func (p *GRPCProvider) SetStateStoreChunkSize(typeName string, size int) {
 	p.mu.Lock()
 	defer p.mu.Unlock()

internal/plugin6/mock_proto/mock.go

@@ -325,6 +325,14 @@ func (s simple) WriteStateBytes(req providers.WriteStateBytesRequest) providers.
 	panic("not implemented")
 }
 
+func (s simple) LockState(req providers.LockStateRequest) providers.LockStateResponse {
+	panic("not implemented")
+}
+
+func (s simple) UnlockState(req providers.UnlockStateRequest) providers.UnlockStateResponse {
+	panic("not implemented")
+}
+
 func (s simple) GetStates(req providers.GetStatesRequest) providers.GetStatesResponse {
 	panic("not implemented")
 }

internal/provider-simple-v6/provider.go

@@ -285,6 +285,14 @@ func (s simple) WriteStateBytes(req providers.WriteStateBytesRequest) providers.
 	panic("not implemented")
 }
 
+func (s simple) LockState(req providers.LockStateRequest) providers.LockStateResponse {
+	panic("not implemented")
+}
+
+func (s simple) UnlockState(req providers.UnlockStateRequest) providers.UnlockStateResponse {
+	panic("not implemented")
+}
+
 func (s simple) GetStates(req providers.GetStatesRequest) providers.GetStatesResponse {
 	// provider-simple uses protocol version 5, which does not include the RPC that maps to this method
 	panic("not implemented")

internal/provider-simple/provider.go

@@ -444,6 +444,14 @@ func (m *Mock) WriteStateBytes(req WriteStateBytesRequest) WriteStateBytesRespon
 	return m.Provider.WriteStateBytes(req)
 }
 
+func (m *Mock) LockState(req LockStateRequest) LockStateResponse {
+	return m.Provider.LockState(req)
+}
+
+func (m *Mock) UnlockState(req UnlockStateRequest) UnlockStateResponse {
+	return m.Provider.UnlockState(req)
+}
+
 func (m *Mock) GetStates(req GetStatesRequest) GetStatesResponse {
 	return m.Provider.GetStates(req)
 }

internal/providers/mock.go

@@ -128,6 +128,11 @@ type Interface interface {
 	// WriteStateBytes streams byte chunks of a given state file into a state store
 	WriteStateBytes(WriteStateBytesRequest) WriteStateBytesResponse
 
+	// LockState locks a given state (i.e. CE workspace)
+	LockState(LockStateRequest) LockStateResponse
+	// UnlockState unlocks a given state (i.e. CE workspace)
+	UnlockState(UnlockStateRequest) UnlockStateResponse
+
 	// GetStates returns a list of all states (i.e. CE workspaces) managed by a given state store
 	GetStates(GetStatesRequest) GetStatesResponse
 	// DeleteState instructs a given state store to delete a specific state (i.e. a CE workspace)
@@ -903,6 +908,29 @@ type WriteStateBytesResponse struct {
 	Diagnostics tfdiags.Diagnostics
 }
 
+type LockStateRequest struct {
+	TypeName  string
+	StateId   string
+	Operation string
+}
+
+type LockStateResponse struct {
+	LockId string
+	// Diagnostics contains any warnings or errors from the method call.
+	Diagnostics tfdiags.Diagnostics
+}
+
+type UnlockStateRequest struct {
+	TypeName string
+	StateId  string
+	LockId   string
+}
+
+type UnlockStateResponse struct {
+	// Diagnostics contains any warnings or errors from the method call.
+	Diagnostics tfdiags.Diagnostics
+}
+
 type GetStatesRequest struct {
 	// TypeName is the name of the state store to request the list of states from
 	TypeName string