/
notebook_instances_fetch.go
78 lines (66 loc) · 2.11 KB
/
notebook_instances_fetch.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package sagemaker
import (
"context"
"github.com/OpsHelmInc/cloudquery/client"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/sagemaker"
"github.com/aws/aws-sdk-go-v2/service/sagemaker/types"
"github.com/cloudquery/plugin-sdk/schema"
)
type WrappedSageMakerNotebookInstance struct {
*sagemaker.DescribeNotebookInstanceOutput
NotebookInstanceArn string
NotebookInstanceName string
}
func fetchSagemakerNotebookInstances(ctx context.Context, meta schema.ClientMeta, _ *schema.Resource, res chan<- any) error {
c := meta.(*client.Client)
svc := c.Services().Sagemaker
config := sagemaker.ListNotebookInstancesInput{}
for {
response, err := svc.ListNotebookInstances(ctx, &config)
if err != nil {
return err
}
res <- response.NotebookInstances
if aws.ToString(response.NextToken) == "" {
break
}
config.NextToken = response.NextToken
}
return nil
}
func getNotebookInstance(ctx context.Context, meta schema.ClientMeta, resource *schema.Resource) error {
c := meta.(*client.Client)
svc := c.Services().Sagemaker
n := resource.Item.(types.NotebookInstanceSummary)
// get more details about the notebook instance
response, err := svc.DescribeNotebookInstance(ctx, &sagemaker.DescribeNotebookInstanceInput{
NotebookInstanceName: n.NotebookInstanceName,
})
if err != nil {
return err
}
resource.Item = &WrappedSageMakerNotebookInstance{
DescribeNotebookInstanceOutput: response,
NotebookInstanceArn: *n.NotebookInstanceArn,
NotebookInstanceName: *n.NotebookInstanceName,
}
return nil
}
func resolveSagemakerNotebookInstanceTags(ctx context.Context, meta schema.ClientMeta, resource *schema.Resource, _ schema.Column) error {
r := resource.Item.(*WrappedSageMakerNotebookInstance)
c := meta.(*client.Client)
svc := c.Services().Sagemaker
config := sagemaker.ListTagsInput{
ResourceArn: &r.NotebookInstanceArn,
}
response, err := svc.ListTags(ctx, &config)
if err != nil {
return err
}
tags := make(map[string]*string, len(response.Tags))
for _, t := range response.Tags {
tags[*t.Key] = t.Value
}
return resource.Set("tags", tags)
}