Skip to content

Commit

Permalink
feat(sagemaker-spaces): add support for Sagemaker Spaces
Browse files Browse the repository at this point in the history
  • Loading branch information
npellegrin committed May 3, 2024
1 parent a51e1a9 commit 16e062b
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 1 deletion.
6 changes: 5 additions & 1 deletion resources/sagemaker-apps.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
})
}
Expand All @@ -70,6 +71,7 @@ type SageMakerApp struct {
appName *string
appType *string
userProfileName *string
spaceName *string
status *string
}

Expand All @@ -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
Expand All @@ -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
}

Expand Down
97 changes: 97 additions & 0 deletions resources/sagemaker-spaces.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 16e062b

Please sign in to comment.