Add snapshot endpoint to management API

iotaledger · Apr 15, 2024 · 57e9e88 · 57e9e88
1 parent 32979ee
commit 57e9e88
Show file tree

Hide file tree

Showing 4 changed files with 232 additions and 1 deletion.
diff --git a/api/management.go b/api/management.go
@@ -56,4 +56,18 @@ type (
 		// The current oldest epoch in the database.
 		Epoch iotago.EpochIndex `serix:""`
 	}
+
+	// CreateSnapshotRequest defines the request of a create snapshot REST API call.
+	CreateSnapshotRequest struct {
+		// The slot of the snapshot.
+		Slot iotago.SlotIndex `serix:""`
+	}
+
+	// CreateSnapshotResponse defines the response of a create snapshot REST API call.
+	CreateSnapshotResponse struct {
+		// The slot of the snapshot.
+		Slot iotago.SlotIndex `serix:""`
+		// The file path of the snapshot file.
+		FilePath string `serix:",lenPrefix=uint8"`
+	}
 )
diff --git a/api/management_test.go b/api/management_test.go
@@ -68,6 +68,21 @@ func Test_ManagementAPIDeSerialize(t *testing.T) {
 			},
 			Target: &api.PruneDatabaseResponse{},
 		},
+		{
+			Name: "ok - CreateSnapshotRequest",
+			Source: &api.CreateSnapshotRequest{
+				Slot: 1,
+			},
+			Target: &api.CreateSnapshotRequest{},
+		},
+		{
+			Name: "ok - CreateSnapshotResponse",
+			Source: &api.CreateSnapshotResponse{
+				Slot:     1,
+				FilePath: "filePath",
+			},
+			Target: &api.CreateSnapshotResponse{},
+		},
 	}
 
 	for _, tt := range tests {
@@ -170,6 +185,26 @@ func Test_ManagementAPIJSONSerialization(t *testing.T) {
 			},
 			Target: `{
 	"epoch": 1
+}`,
+		},
+		{
+			Name: "ok - CreateSnapshotRequest",
+			Source: &api.CreateSnapshotRequest{
+				Slot: 1,
+			},
+			Target: `{
+	"slot": 1
+}`,
+		},
+		{
+			Name: "ok - CreateSnapshotResponse",
+			Source: &api.CreateSnapshotResponse{
+				Slot:     1,
+				FilePath: "filePath",
+			},
+			Target: `{
+	"slot": 1,
+	"filePath": "filePath"
 }`,
 		},
 	}

diff --git a/nodeclient/management_client.go b/nodeclient/management_client.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 
+	iotago "github.com/iotaledger/iota.go/v4"
 	"github.com/iotaledger/iota.go/v4/api"
 )
 
@@ -18,6 +19,14 @@ type (
 		Peers(ctx context.Context) (*api.PeersResponse, error)
 		// AddPeer adds a new peer by libp2p multi address with optional alias.
 		AddPeer(ctx context.Context, multiAddress string, alias ...string) (*api.PeerInfo, error)
+		// PruneDatabaseBySize prunes the database by target size.
+		PruneDatabaseBySize(ctx context.Context, targetDatabaseSize string) (*api.PruneDatabaseResponse, error)
+		// PruneDatabaseByEpoch prunes the database by epoch.
+		PruneDatabaseByEpoch(ctx context.Context, epoch iotago.EpochIndex) (*api.PruneDatabaseResponse, error)
+		// PruneDatabaseByDepth prunes the database by depth.
+		PruneDatabaseByDepth(ctx context.Context, depth iotago.EpochIndex) (*api.PruneDatabaseResponse, error)
+		// CreateSnapshot creates a snapshot.
+		CreateSnapshot(ctx context.Context, slot iotago.SlotIndex) (*api.CreateSnapshotResponse, error)
 	}
 
 	managementClient struct {
@@ -91,3 +100,63 @@ func (client *managementClient) AddPeer(ctx context.Context, multiAddress string
 
 	return res, nil
 }
+
+// PruneDatabaseBySize prunes the database by target size.
+func (client *managementClient) PruneDatabaseBySize(ctx context.Context, targetDatabaseSize string) (*api.PruneDatabaseResponse, error) {
+	req := &api.PruneDatabaseRequest{
+		TargetDatabaseSize: targetDatabaseSize,
+	}
+
+	res := new(api.PruneDatabaseResponse)
+	//nolint:bodyclose
+	if _, err := client.DoWithRequestHeaderHook(ctx, http.MethodPost, api.ManagementRouteDatabasePrune, RequestHeaderHookAcceptJSON, req, res); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// PruneDatabaseByEpoch prunes the database by epoch.
+func (client *managementClient) PruneDatabaseByEpoch(ctx context.Context, epoch iotago.EpochIndex) (*api.PruneDatabaseResponse, error) {
+	req := &api.PruneDatabaseRequest{
+		Epoch: epoch,
+	}
+
+	res := new(api.PruneDatabaseResponse)
+	//nolint:bodyclose
+	if _, err := client.DoWithRequestHeaderHook(ctx, http.MethodPost, api.ManagementRouteDatabasePrune, RequestHeaderHookAcceptJSON, req, res); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// PruneDatabaseByDepth prunes the database by depth.
+func (client *managementClient) PruneDatabaseByDepth(ctx context.Context, depth iotago.EpochIndex) (*api.PruneDatabaseResponse, error) {
+	req := &api.PruneDatabaseRequest{
+		Depth: depth,
+	}
+
+	res := new(api.PruneDatabaseResponse)
+	//nolint:bodyclose
+	if _, err := client.DoWithRequestHeaderHook(ctx, http.MethodPost, api.ManagementRouteDatabasePrune, RequestHeaderHookAcceptJSON, req, res); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// CreateSnapshot creates a snapshot.
+func (client *managementClient) CreateSnapshot(ctx context.Context, slot iotago.SlotIndex) (*api.CreateSnapshotResponse, error) {
+	req := &api.CreateSnapshotRequest{
+		Slot: slot,
+	}
+
+	res := new(api.CreateSnapshotResponse)
+	//nolint:bodyclose
+	if _, err := client.DoWithRequestHeaderHook(ctx, http.MethodPost, api.ManagementRouteSnapshotsCreate, RequestHeaderHookAcceptJSON, req, res); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
diff --git a/nodeclient/management_client_test.go b/nodeclient/management_client_test.go
@@ -168,7 +168,7 @@ func TestManagementClient_AddPeer(t *testing.T) {
 	}
 
 	mockGetJSON(api.RouteRoutes, 200, originRoutes)
-	mockPostJSON(api.ManagementRoutePeers, 201, req, originRes)
+	mockPostJSON(api.ManagementRoutePeers, 200, req, originRes)
 
 	client := nodeClient(t)
 
@@ -179,3 +179,116 @@ func TestManagementClient_AddPeer(t *testing.T) {
 	require.NoError(t, err)
 	require.EqualValues(t, originRes, resp)
 }
+
+func TestManagementClient_PruneDatabaseBySize(t *testing.T) {
+	defer gock.Off()
+
+	targetSize := "1GB"
+
+	originRes := &api.PruneDatabaseResponse{
+		Epoch: 1,
+	}
+
+	req := &api.PruneDatabaseRequest{TargetDatabaseSize: targetSize}
+
+	originRoutes := &api.RoutesResponse{
+		Routes: []iotago.PrefixedStringUint8{api.ManagementPluginName},
+	}
+
+	mockGetJSON(api.RouteRoutes, 200, originRoutes)
+	mockPostJSON(api.ManagementRouteDatabasePrune, 200, req, originRes)
+
+	client := nodeClient(t)
+
+	management, err := client.Management(context.TODO())
+	require.NoError(t, err)
+
+	resp, err := management.PruneDatabaseBySize(context.Background(), targetSize)
+	require.NoError(t, err)
+	require.EqualValues(t, originRes, resp)
+}
+
+func TestManagementClient_PruneDatabaseByEpoch(t *testing.T) {
+	defer gock.Off()
+
+	epoch := iotago.EpochIndex(1)
+
+	originRes := &api.PruneDatabaseResponse{
+		Epoch: 1,
+	}
+
+	req := &api.PruneDatabaseRequest{Epoch: epoch}
+
+	originRoutes := &api.RoutesResponse{
+		Routes: []iotago.PrefixedStringUint8{api.ManagementPluginName},
+	}
+
+	mockGetJSON(api.RouteRoutes, 200, originRoutes)
+	mockPostJSON(api.ManagementRouteDatabasePrune, 200, req, originRes)
+
+	client := nodeClient(t)
+
+	management, err := client.Management(context.TODO())
+	require.NoError(t, err)
+
+	resp, err := management.PruneDatabaseByEpoch(context.Background(), epoch)
+	require.NoError(t, err)
+	require.EqualValues(t, originRes, resp)
+}
+
+func TestManagementClient_PruneDatabaseByDepth(t *testing.T) {
+	defer gock.Off()
+
+	depth := iotago.EpochIndex(1)
+
+	originRes := &api.PruneDatabaseResponse{
+		Epoch: 1,
+	}
+
+	req := &api.PruneDatabaseRequest{Depth: depth}
+
+	originRoutes := &api.RoutesResponse{
+		Routes: []iotago.PrefixedStringUint8{api.ManagementPluginName},
+	}
+
+	mockGetJSON(api.RouteRoutes, 200, originRoutes)
+	mockPostJSON(api.ManagementRouteDatabasePrune, 200, req, originRes)
+
+	client := nodeClient(t)
+
+	management, err := client.Management(context.TODO())
+	require.NoError(t, err)
+
+	resp, err := management.PruneDatabaseByDepth(context.Background(), depth)
+	require.NoError(t, err)
+	require.EqualValues(t, originRes, resp)
+}
+
+func TestManagementClient_CreateSnapshot(t *testing.T) {
+	defer gock.Off()
+
+	slot := iotago.SlotIndex(1)
+
+	originRes := &api.CreateSnapshotResponse{
+		Slot:     1,
+		FilePath: "filePath",
+	}
+
+	req := &api.CreateSnapshotRequest{Slot: slot}
+
+	originRoutes := &api.RoutesResponse{
+		Routes: []iotago.PrefixedStringUint8{api.ManagementPluginName},
+	}
+
+	mockGetJSON(api.RouteRoutes, 200, originRoutes)
+	mockPostJSON(api.ManagementRouteSnapshotsCreate, 200, req, originRes)
+
+	client := nodeClient(t)
+
+	management, err := client.Management(context.TODO())
+	require.NoError(t, err)
+
+	resp, err := management.CreateSnapshot(context.Background(), slot)
+	require.NoError(t, err)
+	require.EqualValues(t, originRes, resp)
+}