Add context and timeout support

cloudquery · Mar 23, 2021 · 5b474c3 · 5b474c3
1 parent e97d69b
commit 5b474c3
Show file tree

Hide file tree

Showing 67 changed files with 143 additions and 156 deletions.
diff --git a/provider/provider.go b/provider/provider.go
@@ -60,6 +60,8 @@ type Config struct {
 	AWSDebug   bool 	 `yaml:"aws_debug"`
 	MaxRetries int       `yaml:"max_retries" default:"5"`
 	MaxBackoff int       `yaml:"max_backoff" default:"30"`
+	// context timeout in seconds
+	Timeout    int       `yaml:"timeout" default:"30"`
 	Resources  []struct {
 		Name  string
 		Other map[string]interface{} `yaml:",inline"`
@@ -255,7 +257,9 @@ func (p *Provider) fetchAccount(accountID string, awsCfg aws.Config, svc *sts.Cl
 				globalServicesFetched[r.Name] = true
 			}
 			g.Go(func() error {
-				err := resourceClients[serviceName].CollectResource(resourceName, resourceConfig)
+				ctx, cancel := context.WithTimeout(ctx, time.Duration(p.config.Timeout)*time.Second)
+				defer cancel()
+				err := resourceClients[serviceName].CollectResource(ctx, resourceName, resourceConfig)
 				if err != nil {
 					var ae smithy.APIError
 					if errors.As(err, &ae) {
@@ -296,6 +300,7 @@ func (p *Provider) Fetch(data []byte) error {
 		return err
 	}
 
+
 	p.regions = p.config.Regions
 	if len(p.regions) == 0 {
 		p.regions = allRegions

diff --git a/resources/autoscaling/client.go b/resources/autoscaling/client.go
@@ -1,6 +1,7 @@
 package autoscaling
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
@@ -30,10 +31,10 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "launch_configurations":
-		return c.launchConfigurations(config)
+		return c.launchConfigurations(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource autoscaling.%s", resource)
 	}

diff --git a/resources/autoscaling/launch_configurations.go b/resources/autoscaling/launch_configurations.go
@@ -150,9 +150,8 @@ var LaunchConfigurationTables = []interface{}{
 	&LaunchConfigurationBlockDeviceMapping{},
 }
 
-func (c *Client) launchConfigurations(gConfig interface{}) error {
+func (c *Client) launchConfigurations(ctx context.Context, gConfig interface{}) error {
 	var config autoscaling.DescribeLaunchConfigurationsInput
-	ctx := context.Background()
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {
 		return err

diff --git a/resources/cloudtrail/client.go b/resources/cloudtrail/client.go
@@ -1,6 +1,7 @@
 package cloudtrail
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/cloudtrail"
@@ -30,10 +31,10 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "trails":
-		return c.trails(config)
+		return c.trails(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource autoscaling.%s", resource)
 	}

diff --git a/resources/cloudtrail/trails.go b/resources/cloudtrail/trails.go
@@ -136,9 +136,8 @@ var TrailTables = []interface{}{
 	&EventSelector{},
 }
 
-func (c *Client) trails(gConfig interface{}) error {
+func (c *Client) trails(ctx context.Context, gConfig interface{}) error {
 	var config cloudtrail.DescribeTrailsInput
-	ctx := context.Background()
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {
 		return err

diff --git a/resources/cloudwatch/alarms.go b/resources/cloudwatch/alarms.go
@@ -163,8 +163,7 @@ var MetricAlarmTables = []interface{} {
 	&MetricAlarmMetric{},
 }
 
-func (c *Client)alarms(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client)alarms(ctx context.Context, gConfig interface{}) error {
 	var config cloudwatch.DescribeAlarmsInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/cloudwatch/client.go b/resources/cloudwatch/client.go
@@ -1,6 +1,7 @@
 package cloudwatch
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/cloudwatch"
@@ -30,10 +31,10 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "alarms":
-		return c.alarms(config)
+		return c.alarms(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource autoscaling.%s", resource)
 	}

diff --git a/resources/cloudwatchlogs/client.go b/resources/cloudwatchlogs/client.go
@@ -1,6 +1,7 @@
 package cloudwatchlogs
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs"
@@ -30,10 +31,10 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "metric_filters":
-		return c.metricFilters(config)
+		return c.metricFilters(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource autoscaling.%s", resource)
 	}

diff --git a/resources/cloudwatchlogs/metric_filters.go b/resources/cloudwatchlogs/metric_filters.go
@@ -80,8 +80,7 @@ var MetricFilterTables = []interface{} {
 	&MetricFilterTransformations{},
 }
 
-func (c *Client)metricFilters(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client)metricFilters(ctx context.Context, gConfig interface{}) error {
 	var config cloudwatchlogs.DescribeMetricFiltersInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/directconnect/client.go b/resources/directconnect/client.go
@@ -1,6 +1,7 @@
 package directconnect
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/directconnect"
@@ -30,10 +31,10 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "gateways":
-		return c.gateways(config)
+		return c.gateways(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource directconnect.%s", resource)
 	}

diff --git a/resources/directconnect/gateways.go b/resources/directconnect/gateways.go
@@ -46,8 +46,7 @@ var GatewayTables = []interface{}{
 	&Gateway{},
 }
 
-func (c *Client) gateways(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) gateways(ctx context.Context, gConfig interface{}) error {
 	var config directconnect.DescribeDirectConnectGatewaysInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/byoip_cidr.go b/resources/ec2/byoip_cidr.go
@@ -41,8 +41,7 @@ var ByoipCidrTables = []interface{}{
 	&ByoipCidr{},
 }
 
-func (c *Client) byoipCidrs(_ interface{}) error {
-	ctx := context.Background()
+func (c *Client) byoipCidrs(ctx context.Context, _ interface{}) error {
 	config := ec2.DescribeByoipCidrsInput{
 		MaxResults: 100,
 	}

diff --git a/resources/ec2/client.go b/resources/ec2/client.go
@@ -1,6 +1,7 @@
 package ec2
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
@@ -30,34 +31,34 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "images":
-		return c.images(config)
+		return c.images(ctx, config)
 	case "instances":
-		return c.instances(config)
+		return c.instances(ctx, config)
 	case "byoip_cidrs":
-		return c.byoipCidrs(config)
+		return c.byoipCidrs(ctx, config)
 	case "customer_gateways":
-		return c.customerGateways(config)
+		return c.customerGateways(ctx, config)
 	case "internet_gateways":
-		return c.internetGateways(config)
+		return c.internetGateways(ctx, config)
 	case "nat_gateways":
-		return c.natGateways(config)
+		return c.natGateways(ctx, config)
 	case "network_acls":
-		return c.networkAcls(config)
+		return c.networkAcls(ctx, config)
 	case "route_tables":
-		return c.routeTables(config)
+		return c.routeTables(ctx, config)
 	case "security_groups":
-		return c.securityGroups(config)
+		return c.securityGroups(ctx, config)
 	case "vpcs":
-		return c.vpcs(config)
+		return c.vpcs(ctx, config)
 	case "subnets":
-		return c.subnets(config)
+		return c.subnets(ctx, config)
 	case "flow_logs":
-		return c.FlowLogs(config)
+		return c.FlowLogs(ctx, config)
 	case "vpc_peering_connections":
-		return c.vpcPeeringConnections(config)
+		return c.vpcPeeringConnections(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource ec2.%s", resource)
 	}

diff --git a/resources/ec2/customer_gateways.go b/resources/ec2/customer_gateways.go
@@ -77,8 +77,7 @@ var CustomerGatewayTables = []interface{}{
 	&CustomerGatewayTag{},
 }
 
-func (c *Client) customerGateways(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) customerGateways(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeCustomerGatewaysInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/flow_logs.go b/resources/ec2/flow_logs.go
@@ -92,9 +92,8 @@ var FlowLogsTables = []interface{}{
 	&FlowLogTag{},
 }
 
-func (c *Client) FlowLogs(gConfig interface{}) error {
+func (c *Client) FlowLogs(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeFlowLogsInput
-	ctx := context.Background()
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {
 		return err

diff --git a/resources/ec2/images.go b/resources/ec2/images.go
@@ -203,8 +203,7 @@ var ImageTables = []interface{}{
 	&ImageTag{},
 }
 
-func (c *Client) images(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) images(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeImagesInput
 	err := mapstructure.Decode(gConfig, &config)
 	if config.Owners == nil {

diff --git a/resources/ec2/instances.go b/resources/ec2/instances.go
@@ -574,8 +574,7 @@ var InstanceTables = []interface{}{
 	&InstanceTag{},
 }
 
-func (c *Client) instances(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) instances(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeInstancesInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/internet_gateways.go b/resources/ec2/internet_gateways.go
@@ -100,8 +100,7 @@ var InternetGatewayTables = []interface{}{
 	&InternetGatewayTag{},
 }
 
-func (c *Client) internetGateways(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) internetGateways(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeInternetGatewaysInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/nat_gateways.go b/resources/ec2/nat_gateways.go
@@ -133,8 +133,7 @@ var NatGatewayTables = []interface{}{
 	&NatGatewayTag{},
 }
 
-func (c *Client) natGateways(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) natGateways(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeNatGatewaysInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/network_acls.go b/resources/ec2/network_acls.go
@@ -166,8 +166,7 @@ var NetworkAclTables = []interface{}{
 	&NetworkAclTag{},
 }
 
-func (c *Client) networkAcls(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) networkAcls(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeNetworkAclsInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/route_tabls.go b/resources/ec2/route_tabls.go
@@ -200,8 +200,7 @@ var RouteTableTables = []interface{}{
 	&RouteTableTag{},
 }
 
-func (c *Client) routeTables(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) routeTables(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeRouteTablesInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/security_groups.go b/resources/ec2/security_groups.go
@@ -229,8 +229,7 @@ var SecurityGroupTables = []interface{}{
 	&SecurityGroupTag{},
 }
 
-func (c *Client) securityGroups(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) securityGroups(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeSecurityGroupsInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/subnets.go b/resources/ec2/subnets.go
@@ -131,8 +131,7 @@ var SubnetTables = []interface{}{
 	&SubnetTag{},
 }
 
-func (c *Client) subnets(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) subnets(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeSubnetsInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/vpc_peering_connections.go b/resources/ec2/vpc_peering_connections.go
@@ -241,8 +241,7 @@ var VPCPeeringConnectionTables = []interface{}{
 	&VpcPeeringConnectionTag{},
 }
 
-func (c *Client) vpcPeeringConnections(gConfig interface{}) error {
-	ctx := context.Background()
+func (c *Client) vpcPeeringConnections(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeVpcPeeringConnectionsInput
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {

diff --git a/resources/ec2/vpcs.go b/resources/ec2/vpcs.go
@@ -163,9 +163,8 @@ var VPCTables = []interface{}{
 	&VpcTag{},
 }
 
-func (c *Client) vpcs(gConfig interface{}) error {
+func (c *Client) vpcs(ctx context.Context, gConfig interface{}) error {
 	var config ec2.DescribeVpcsInput
-	ctx := context.Background()
 	err := mapstructure.Decode(gConfig, &config)
 	if err != nil {
 		return err

diff --git a/resources/ecr/client.go b/resources/ecr/client.go
@@ -1,6 +1,7 @@
 package ecr
 
 import (
+	"context"
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/ecr"
@@ -30,10 +31,10 @@ func NewClient(awsConfig aws.Config, db *database.Database, log hclog.Logger,
 	}
 }
 
-func (c *Client) CollectResource(resource string, config interface{}) error {
+func (c *Client) CollectResource(ctx context.Context, resource string, config interface{}) error {
 	switch resource {
 	case "images":
-		return c.images(config)
+		return c.images(ctx, config)
 	default:
 		return fmt.Errorf("unsupported resource ecr.%s", resource)
 	}