Skip to content

Commit

Permalink
Adding rekognition project and dataset support.
Browse files Browse the repository at this point in the history
Rekogniton datasets don't have list functions directly available.  Info has been parsed from the DescribeProjects call.
  • Loading branch information
swhite-oreilly authored and gsoria committed Oct 4, 2023
1 parent 91e2670 commit a90dcf2
Show file tree
Hide file tree
Showing 2 changed files with 122 additions and 0 deletions.
62 changes: 62 additions & 0 deletions resources/rekognition-dataset.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
package resources

import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/rekognition"
)

type RekognitionDataset struct {
svc *rekognition.Rekognition
arn *string
}

func init() {
register("RekognitionDataset", ListRekognitionDatasets)
}

func ListRekognitionDatasets(sess *session.Session) ([]Resource, error) {
svc := rekognition.New(sess)
resources := []Resource{}

params := &rekognition.DescribeProjectsInput{
MaxResults: aws.Int64(100),
}

for {
output, err := svc.DescribeProjects(params)
if err != nil {
return nil, err
}

for _, project := range output.ProjectDescriptions {
for _, dataset := range project.Datasets {
resources = append(resources, &RekognitionDataset{
svc: svc,
arn: dataset.DatasetArn,
})
}
}

if output.NextToken == nil {
break
}

params.NextToken = output.NextToken
}

return resources, nil
}

func (f *RekognitionDataset) Remove() error {

_, err := f.svc.DeleteDataset(&rekognition.DeleteDatasetInput{
DatasetArn: f.arn,
})

return err
}

func (f *RekognitionDataset) String() string {
return *f.arn
}
60 changes: 60 additions & 0 deletions resources/rekognition-project.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
package resources

import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/rekognition"
)

type RekognitionProject struct {
svc *rekognition.Rekognition
arn *string
}

func init() {
register("RekognitionProject", ListRekognitionProjects)
}

func ListRekognitionProjects(sess *session.Session) ([]Resource, error) {
svc := rekognition.New(sess)
resources := []Resource{}

params := &rekognition.DescribeProjectsInput{
MaxResults: aws.Int64(100),
}

for {
output, err := svc.DescribeProjects(params)
if err != nil {
return nil, err
}

for _, project := range output.ProjectDescriptions {
resources = append(resources, &RekognitionProject{
svc: svc,
arn: project.ProjectArn,
})
}

if output.NextToken == nil {
break
}

params.NextToken = output.NextToken
}

return resources, nil
}

func (f *RekognitionProject) Remove() error {

_, err := f.svc.DeleteProject(&rekognition.DeleteProjectInput{
ProjectArn: f.arn,
})

return err
}

func (f *RekognitionProject) String() string {
return *f.arn
}

0 comments on commit a90dcf2

Please sign in to comment.