From 16e062bcff04d837f0092d9d3d1f8afc743d078a Mon Sep 17 00:00:00 2001
From: Nicolas Pellegrin <npellegrin.dev@gmail.com>
Date: Fri, 3 May 2024 11:10:28 +0200
Subject: [PATCH] feat(sagemaker-spaces): add support for Sagemaker Spaces

---
 resources/sagemaker-apps.go   |  6 ++-
 resources/sagemaker-spaces.go | 97 +++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 resources/sagemaker-spaces.go

diff --git a/resources/sagemaker-apps.go b/resources/sagemaker-apps.go
index d510c3fc..9246f147 100644
--- a/resources/sagemaker-apps.go
+++ b/resources/sagemaker-apps.go
@@ -50,6 +50,7 @@ func (l *SageMakerAppLister) List(_ context.Context, o interface{}) ([]resource.
 				appName:         app.AppName,
 				appType:         app.AppType,
 				userProfileName: app.UserProfileName,
+				spaceName:       app.SpaceName,
 				status:          app.Status,
 			})
 		}
@@ -70,6 +71,7 @@ type SageMakerApp struct {
 	appName         *string
 	appType         *string
 	userProfileName *string
+	spaceName       *string
 	status          *string
 }
 
@@ -79,6 +81,7 @@ func (f *SageMakerApp) Remove(_ context.Context) error {
 		AppName:         f.appName,
 		AppType:         f.appType,
 		UserProfileName: f.userProfileName,
+		SpaceName:       f.spaceName,
 	})
 
 	return err
@@ -94,7 +97,8 @@ func (f *SageMakerApp) Properties() types.Properties {
 		Set("DomainID", f.domainID).
 		Set("AppName", f.appName).
 		Set("AppType", f.appType).
-		Set("UserProfileName", f.userProfileName)
+		Set("UserProfileName", f.userProfileName).
+		Set("SpaceName", f.spaceName)
 	return properties
 }
 
diff --git a/resources/sagemaker-spaces.go b/resources/sagemaker-spaces.go
new file mode 100644
index 00000000..69bd3269
--- /dev/null
+++ b/resources/sagemaker-spaces.go
@@ -0,0 +1,97 @@
+package resources
+
+import (
+	"context"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/sagemaker"
+
+	"github.com/ekristen/libnuke/pkg/registry"
+	"github.com/ekristen/libnuke/pkg/resource"
+	"github.com/ekristen/libnuke/pkg/types"
+
+	"github.com/ekristen/aws-nuke/pkg/nuke"
+)
+
+const SageMakerSpaceResource = "SageMakerSpace"
+
+func init() {
+	registry.Register(&registry.Registration{
+		Name:   SageMakerSpaceResource,
+		Scope:  nuke.Account,
+		Lister: &SageMakerSpaceLister{},
+	})
+}
+
+type SageMakerSpaceLister struct{}
+
+func (l *SageMakerSpaceLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
+	opts := o.(*nuke.ListerOpts)
+
+	svc := sagemaker.New(opts.Session)
+	resources := make([]resource.Resource, 0)
+
+	params := &sagemaker.ListSpacesInput{
+		MaxResults: aws.Int64(30),
+	}
+
+	for {
+		resp, err := svc.ListSpaces(params)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, space := range resp.Spaces {
+			resources = append(resources, &SageMakerSpace{
+				svc:              svc,
+				domainID:         space.DomainId,
+				spaceDisplayName: space.SpaceDisplayName,
+				spaceName:        space.SpaceName,
+				status:           space.Status,
+				lastModifiedTime: space.LastModifiedTime,
+			})
+		}
+
+		if resp.NextToken == nil {
+			break
+		}
+
+		params.NextToken = resp.NextToken
+	}
+
+	return resources, nil
+}
+
+type SageMakerSpace struct {
+	svc              *sagemaker.SageMaker
+	domainID         *string
+	spaceDisplayName *string
+	spaceName        *string
+	status           *string
+	lastModifiedTime *time.Time
+}
+
+func (f *SageMakerSpace) Remove(_ context.Context) error {
+	_, err := f.svc.DeleteSpace(&sagemaker.DeleteSpaceInput{
+		DomainId:  f.domainID,
+		SpaceName: f.spaceName,
+	})
+
+	return err
+}
+
+func (f *SageMakerSpace) String() string {
+	return *f.spaceName
+}
+
+func (f *SageMakerSpace) Properties() types.Properties {
+	properties := types.NewProperties()
+	properties.
+		Set("DomainID", f.domainID).
+		Set("SpaceDisplayName", f.spaceDisplayName).
+		Set("SpaceName", f.spaceName).
+		Set("Status", f.status).
+		Set("LastModifiedTime", f.lastModifiedTime)
+	return properties
+}