Update aws_ec2_instance, aws_iam_policy, aws_rds_db_cluster_snapshot for better use of API filters, context cancellation in list calls and page limiting for limit clause in query

aws/table_aws_ec2_instance.go

@@ -28,6 +28,26 @@ func tableAwsEc2Instance(_ context.Context) *plugin.Table {
 		},
 		List: &plugin.ListConfig{
 			Hydrate: listEc2Instance,
+			KeyColumns: []*plugin.KeyColumn{
+				{Name: "hypervisor", Require: plugin.Optional},
+				{Name: "iam_instance_profile_arn", Require: plugin.Optional},
+				{Name: "image_id", Require: plugin.Optional},
+				{Name: "instance_lifecycle", Require: plugin.Optional},
+				{Name: "instance_state", Require: plugin.Optional},
+				{Name: "instance_type", Require: plugin.Optional},
+				{Name: "monitoring_state", Require: plugin.Optional},
+				{Name: "outpost_arn", Require: plugin.Optional},
+				{Name: "placement_availability_zone", Require: plugin.Optional},
+				{Name: "placement_group_name", Require: plugin.Optional},
+				{Name: "public_dns_name", Require: plugin.Optional},
+				{Name: "ram_disk_id", Require: plugin.Optional},
+				{Name: "root_device_name", Require: plugin.Optional},
+				{Name: "root_device_type", Require: plugin.Optional},
+				{Name: "subnet_id", Require: plugin.Optional},
+				{Name: "placement_tenancy", Require: plugin.Optional},
+				{Name: "virtualization_type", Require: plugin.Optional},
+				{Name: "vpc_id", Require: plugin.Optional},
+			},
 		},
 		GetMatrixItem: BuildRegionList,
 		Columns: awsRegionalColumns([]*plugin.Column{
@@ -327,19 +347,48 @@ func listEc2Instance(ctx context.Context, d *plugin.QueryData, _ *plugin.Hydrate
 		return nil, err
 	}
 
+	input := ec2.DescribeInstancesInput{
+		MaxResults: types.Int64(1000),
+	}
+	filters := buildEc2InstanceFilter(d.KeyColumnQuals)
+
+	if len(filters) != 0 {
+		input.Filters = filters
+	}
+
+	// If the requested number of items is less than the paging max limit
+	// set the limit to that instead
+	limit := d.QueryContext.Limit
+	if d.QueryContext.Limit != nil {
+		if *limit < *input.MaxResults {
+			// select * from aws_ec2_instance limit 1
+			// Error: InvalidParameterValue: Value ( 1 ) for parameter maxResults is invalid. Expecting a value greater than 5.
+			// 		status code: 400, request id: a84912d9-f5fd-403f-8e37-7f7b3f6faba6
+			if *limit < 5 {
+				input.MaxResults = types.Int64(5)
+			} else {
+				input.MaxResults = limit
+			}
+		}
+	}
+
 	// List call
-	err = svc.DescribeInstancesPages(
-		&ec2.DescribeInstancesInput{},
-		func(page *ec2.DescribeInstancesOutput, isLast bool) bool {
-			if page.Reservations != nil && len(page.Reservations) > 0 {
-				for _, reservation := range page.Reservations {
-					for _, instance := range reservation.Instances {
-						d.StreamListItem(ctx, instance)
+	err = svc.DescribeInstancesPages(&input, func(page *ec2.DescribeInstancesOutput, isLast bool) bool {
+		if page.Reservations != nil && len(page.Reservations) > 0 {
+			for _, reservation := range page.Reservations {
+				for _, instance := range reservation.Instances {
+					d.StreamListItem(ctx, instance)
+					// Check if context has been cancelled or if the limit has been hit (if specified)
+					// if there is a limit, it will return the number of rows required to reach this limit
+					if d.QueryStatus.RowsRemaining(ctx) == 0 {
+						return true
 					}
 				}
 			}
-			return !isLast
-		},
+		}
+
+		return !isLast
+	},
 	)
 
 	if err != nil {
@@ -610,3 +659,55 @@ func ec2InstanceStateChangeTime(_ context.Context, d *transform.TransformData) (
 	}
 	return data.LaunchTime, nil
 }
+
+//// UTILITY FUNCTIONS
+
+// build ec2 instance list call input filter
+func buildEc2InstanceFilter(equalQuals plugin.KeyColumnEqualsQualMap) []*ec2.Filter {
+	filters := make([]*ec2.Filter, 0)
+
+	filterQuals := map[string]string{
+		"hypervisor":                  "hypervisor",
+		"iam_instance_profile_arn":    "iam-instance-profile.arn",
+		"image_id":                    "image-id",
+		"instance_lifecycle":          "instance-lifecycle",
+		"instance_state":              "instance-state-name",
+		"instance_type":               "instance-type",
+		"monitoring_state":            "monitoring-state",
+		"outpost_arn":                 "outpost-arn",
+		"placement_availability_zone": "availability-zone",
+		"placement_group_name":        "placement-group-name",
+		"public_dns_name":             "dns-name",
+		"ram_disk_id":                 "ramdisk-id",
+		"root_device_name":            "root-device-name",
+		"root_device_type":            "root-device-type",
+		"subnet_id":                   "subnet-id",
+		"placement_tenancy":           "tenancy",
+		"virtualization_type":         "virtualization-type",
+		"vpc_id":                      "vpc-id",
+	}
+
+	for columnName, filterName := range filterQuals {
+		if equalQuals[columnName] != nil {
+			filter := ec2.Filter{
+				Name: types.String(filterName),
+			}
+			value := equalQuals[columnName]
+			if value.GetStringValue() != "" {
+				filter.Values = []*string{types.String(equalQuals[columnName].GetStringValue())}
+			} else if value.GetListValue() != nil {
+				filter.Values = getListValues(value.GetListValue())
+			}
+			filters = append(filters, &filter)
+		}
+	}
+	return filters
+}
+
+func getListValues(listValue *proto.QualValueList) []*string {
+	values := make([]*string, 0)
+	for _, value := range listValue.Values {
+		values = append(values, types.String(value.GetStringValue()))
+	}
+	return values
+}

aws/table_aws_iam_policy.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"strings"
 
+	"github.com/turbot/go-kit/types"
 	"github.com/turbot/steampipe-plugin-sdk/grpc/proto"
 	"github.com/turbot/steampipe-plugin-sdk/plugin/transform"
 
@@ -23,6 +24,11 @@ func tableAwsIamPolicy(_ context.Context) *plugin.Table {
 		},
 		List: &plugin.ListConfig{
 			Hydrate: listIamPolicies,
+			KeyColumns: plugin.KeyColumnSlice{
+				{Name: "is_aws_managed", Require: plugin.Optional, Operators: []string{"<>", "="}},
+				{Name: "is_attached", Require: plugin.Optional, Operators: []string{"<>", "="}},
+				{Name: "path", Require: plugin.Optional},
+			},
 		},
 		Columns: awsColumns([]*plugin.Column{
 			{
@@ -73,6 +79,12 @@ func tableAwsIamPolicy(_ context.Context) *plugin.Table {
 				Description: "The number of entities (users, groups, and roles) that the policy is attached to.",
 				Type:        proto.ColumnType_INT,
 			},
+			{
+				Name:        "is_attached",
+				Description: "Specifies whether the policy is attached to at least one IAM user, group, or role.",
+				Type:        proto.ColumnType_BOOL,
+				Transform:   transform.FromField("AttachmentCount").Transform(attachementCountToBool),
+			},
 			{
 				Name:        "default_version_id",
 				Description: "The identifier for the version of the policy that is set as the default version.",
@@ -138,14 +150,31 @@ func listIamPolicies(ctx context.Context, d *plugin.QueryData, _ *plugin.Hydrate
 		return nil, err
 	}
 
-	err = svc.ListPoliciesPages(
-		&iam.ListPoliciesInput{},
-		func(page *iam.ListPoliciesOutput, lastPage bool) bool {
-			for _, policy := range page.Policies {
-				d.StreamListItem(ctx, policy)
+	input := buildIamPolicyFilter(d.KeyColumnQuals, d.Quals)
+	input.MaxItems = types.Int64(100)
+
+	// If the requested number of items is less than the paging max limit
+	// set the limit to that instead
+	limit := d.QueryContext.Limit
+	if d.QueryContext.Limit != nil {
+		if *limit < *input.MaxItems {
+			input.MaxItems = limit
+		}
+	}
+
+	// List call
+	err = svc.ListPoliciesPages(&input, func(page *iam.ListPoliciesOutput, lastPage bool) bool {
+		for _, policy := range page.Policies {
+			d.StreamListItem(ctx, policy)
+			// Check if context has been cancelled or if the limit has been hit (if specified)
+			// if there is a limit, it will return the number of rows required to reach this limit
+			if d.QueryStatus.RowsRemaining(ctx) == 0 {
+				return true
 			}
-			return !lastPage
-		},
+		}
+
+		return !lastPage
+	},
 	)
 	return nil, err
 }
@@ -229,7 +258,7 @@ func isPolicyAwsManaged(ctx context.Context, d *plugin.QueryData, h *plugin.Hydr
 	return false, nil
 }
 
-//// Transform Function
+//// TRANSFORM FUNCTIONS
 
 func iamPolicyTurbotTags(_ context.Context, d *transform.TransformData) (interface{}, error) {
 	policy := d.HydrateItem.(*iam.Policy)
@@ -245,3 +274,82 @@ func iamPolicyTurbotTags(_ context.Context, d *transform.TransformData) (interfa
 
 	return &turbotTagsMap, nil
 }
+
+func attachementCountToBool(_ context.Context, d *transform.TransformData) (interface{}, error) {
+	attachementCount := types.Int64Value((d.Value.(*int64)))
+	if attachementCount == 0 {
+		return false, nil
+	}
+	return true, nil
+}
+
+//// UTILITY FUNCTIONS
+
+func buildIamPolicyFilter(equalQuals plugin.KeyColumnEqualsQualMap, quals plugin.KeyColumnQualMap) iam.ListPoliciesInput {
+	input := iam.ListPoliciesInput{}
+
+	filterQuals := FilterQuals{
+		FilterQual{"is_aws_managed", "Bool", "Scope"},
+		FilterQual{"is_attached", "Bool", "OnlyAttached"},
+		FilterQual{"path", "String", "PathPrefix"},
+	}
+
+	// EqualsQualMap handling
+	for _, filterQual := range filterQuals {
+		if equalQuals[filterQual.ColumnName] != nil {
+			switch filterQual.ColumnType {
+			case "String":
+				input.PathPrefix = types.String(equalQuals[filterQual.ColumnName].GetStringValue())
+			case "Bool":
+				if filterQual.ColumnName == "is_aws_managed" {
+					input.SetScope("Local")
+					if equalQuals[filterQual.ColumnName].GetBoolValue() {
+						input.SetScope("AWS")
+					}
+				}
+				if filterQual.ColumnName == "is_attached" {
+					input.SetOnlyAttached(false)
+					if equalQuals[filterQual.ColumnName].GetBoolValue() {
+						input.SetOnlyAttached(true)
+					}
+				}
+			}
+		}
+	}
+
+	boolNEQuals := []string{
+		"is_aws_managed",
+		"is_attached",
+	}
+	// Non-Equals Qual Map handling
+	for _, qual := range boolNEQuals {
+		if quals[qual] != nil {
+			for _, q := range quals[qual].Quals {
+				value := q.Value.GetBoolValue()
+				if q.Operator == "<>" {
+					if qual == "is_aws_managed" {
+						input.SetScope("Local")
+						if !value {
+							input.SetScope("AWS")
+						}
+					}
+					if qual == "is_attached" {
+						input.SetOnlyAttached(false)
+						if !value {
+							input.SetOnlyAttached(true)
+						}
+					}
+				}
+			}
+		}
+	}
+	return input
+}
+
+type FilterQuals []FilterQual
+
+type FilterQual struct {
+	ColumnName   string
+	ColumnType   string
+	PropertyName string
+}