[aws][fix] S3 metrics for all storage types (#2058)

Author: aquamatthias

plugins/aws/fix_plugin_aws/resource/cloudwatch.py

@@ -418,6 +418,7 @@ class AwsCloudwatchQuery:
     metric_id: str
     stat: str = "Sum"
     unit: str = "Count"
+    fix_metric_name: Optional[str] = None  # Override the default metric name. The name is taken AS IS.
 
     def to_json(self) -> Json:
         return {
@@ -437,13 +438,15 @@ def to_json(self) -> Json:
 
     @staticmethod
     def create(
+        *,
         metric_name: str,
         namespace: str,
         period: timedelta,
         ref_id: str,
         metric_id: Optional[str] = None,
         stat: str = "Sum",
         unit: str = "Count",
+        fix_metric_name: Optional[str] = None,
         **dimensions: str,
     ) -> "AwsCloudwatchQuery":
         dims = "_".join(f"{k}+{v}" for k, v in dimensions.items())
@@ -458,6 +461,7 @@ def create(
             metric_id=rid,
             stat=stat,
             unit=unit,
+            fix_metric_name=fix_metric_name,
         )
 
 
@@ -571,7 +575,8 @@ def update_resource_metrics(
 
         for metric_value, maybe_stat_name in normalizer.compute_stats(metric.metric_values):
             try:
-                name = normalizer.metric_name + "_" + normalizer.unit
+                metric_name = query.fix_metric_name or normalizer.metric_name
+                name = metric_name + "_" + normalizer.unit
                 value = normalizer.normalize_value(metric_value)
                 stat_name = str(maybe_stat_name or normalizer.stat_map[query.stat])
 

plugins/aws/fix_plugin_aws/resource/ec2.py

@@ -572,8 +572,16 @@ def update_atime_mtime() -> None:
                 elif volume.volume_status == VolumeStatus.AVAILABLE:
                     vid = volume.id
                     lookup[vid] = volume
-                    queries.append(AwsCloudwatchQuery.create("VolumeReadOps", "AWS/EBS", delta, vid, VolumeId=vid))
-                    queries.append(AwsCloudwatchQuery.create("VolumeWriteOps", "AWS/EBS", delta, vid, VolumeId=vid))
+                    queries.append(
+                        AwsCloudwatchQuery.create(
+                            metric_name="VolumeReadOps", namespace="AWS/EBS", period=delta, ref_id=vid, VolumeId=vid
+                        )
+                    )
+                    queries.append(
+                        AwsCloudwatchQuery.create(
+                            metric_name="VolumeWriteOps", namespace="AWS/EBS", period=delta, ref_id=vid, VolumeId=vid
+                        )
+                    )
 
             for query, metric in AwsCloudwatchMetricData.query_for(builder, queries, start, now).items():
                 if non_zero := metric.first_non_zero():

plugins/aws/fix_plugin_aws/resource/rds.py

@@ -507,7 +507,13 @@ def update_atime_mtime() -> None:
                 vid = rds.id
                 lookup[vid] = rds
                 queries.append(
-                    AwsCloudwatchQuery.create("DatabaseConnections", "AWS/RDS", delta, vid, DBInstanceIdentifier=vid)
+                    AwsCloudwatchQuery.create(
+                        metric_name="DatabaseConnections",
+                        namespace="AWS/RDS",
+                        period=delta,
+                        ref_id=vid,
+                        DBInstanceIdentifier=vid,
+                    )
                 )
 
             for query, metric in AwsCloudwatchMetricData.query_for(builder, queries, start, now).items():

plugins/aws/fix_plugin_aws/resource/s3.py

@@ -1,4 +1,5 @@
 import logging
+from collections import defaultdict
 from json import loads as json_loads
 from typing import ClassVar, Dict, List, Type, Optional, cast, Any
 
@@ -343,6 +344,14 @@ def _get_tags(self, client: AwsClient) -> Dict[str, str]:
 
     @classmethod
     def collect_usage_metrics(cls: Type[AwsResource], builder: GraphBuilder) -> None:
+        storage_types = {
+            "StandardStorage": "standard_storage",
+            "IntelligentTieringStorage": "intelligent_tiering_storage",
+            "StandardIAStorage": "standard_ia_storage",
+            "OneZoneIAStorage": "one_zone_ia_storage",
+            "GlacierStorage": "glacier_storage",
+            "DeepArchiveStorage": "deep_archive_storage",
+        }
         for region in {
             s3_bucket.bucket_location
             for s3_bucket in builder.nodes(clazz=AwsS3Bucket)
@@ -367,18 +376,20 @@ def collect_usage_metrics(cls: Type[AwsResource], builder: GraphBuilder) -> None
                         StorageType="AllStorageTypes",
                     )
                 )
-                queries.append(
-                    AwsCloudwatchQuery.create(
-                        metric_name="BucketSizeBytes",
-                        namespace="AWS/S3",
-                        period=delta,
-                        ref_id=s3_id,
-                        stat="Average",
-                        unit="Bytes",
-                        BucketName=s3.name or s3.safe_name,
-                        StorageType="StandardStorage",
+                for storage_type, storage_type_name in storage_types.items():
+                    queries.append(
+                        AwsCloudwatchQuery.create(
+                            metric_name="BucketSizeBytes",
+                            namespace="AWS/S3",
+                            period=delta,
+                            ref_id=s3_id,
+                            stat="Average",
+                            unit="Bytes",
+                            fix_metric_name=f"{storage_type_name}_bucket_size",
+                            BucketName=s3.name or s3.safe_name,
+                            StorageType=storage_type,
+                        )
                     )
-                )
             metric_normalizers = {
                 "BucketSizeBytes": MetricNormalization(
                     metric_name=MetricName.BucketSizeBytes,
@@ -395,6 +406,15 @@ def collect_usage_metrics(cls: Type[AwsResource], builder: GraphBuilder) -> None
             region_builder = builder.for_region(AwsRegion(id=region, name=region))
             cloudwatch_result = AwsCloudwatchMetricData.query_for(region_builder, queries, start, now)
             update_resource_metrics(s3s, cloudwatch_result, metric_normalizers)
+            # Calculate the total bucket size for each bucket by summing up the sizes of all storage types
+            for s3 in s3s.values():
+                bucket_size: Dict[str, float] = defaultdict(float)
+                for metric_name, metric_values in s3._resource_usage.items():
+                    if metric_name.endswith("_bucket_size_bytes"):
+                        for name, value in metric_values.items():
+                            bucket_size[name] += value
+                if bucket_size:
+                    s3._resource_usage["bucket_size_bytes"] = dict(bucket_size)
 
     def update_resource_tag(self, client: AwsClient, key: str, value: str) -> bool:
         tags = self._get_tags(client)

plugins/aws/test/resources/cloudwatch_test.py

@@ -39,8 +39,20 @@ def test_metrics_filter() -> None:
 def test_metric(builder: GraphBuilder) -> None:
     now = datetime(2020, 3, 1, tzinfo=timezone.utc)
     earlier = now - timedelta(days=60)
-    read = AwsCloudwatchQuery.create("VolumeReadOps", "AWS/EBS", timedelta(hours=1), "vol-123", VolumeId="vol-123")
-    write = AwsCloudwatchQuery.create("VolumeWriteOps", "AWS/EBS", timedelta(hours=1), "vol-123", VolumeId="vol-123")
+    read = AwsCloudwatchQuery.create(
+        metric_name="VolumeReadOps",
+        namespace="AWS/EBS",
+        period=timedelta(hours=1),
+        ref_id="vol-123",
+        VolumeId="vol-123",
+    )
+    write = AwsCloudwatchQuery.create(
+        metric_name="VolumeWriteOps",
+        namespace="AWS/EBS",
+        period=timedelta(hours=1),
+        ref_id="vol-123",
+        VolumeId="vol-123",
+    )
     result = AwsCloudwatchMetricData.query_for(builder, [read, write], earlier, now)
     assert result[read].first_non_zero() == (datetime(2020, 1, 18, 17, 40, tzinfo=timezone.utc), 15.0)
     assert result[write].first_non_zero() == (datetime(2020, 1, 18, 18, 40, tzinfo=timezone.utc), 12861.0)