Add groupby feature

pyproject.toml

@@ -10,7 +10,7 @@ keywords = [
 license = "Apache 2.0"
 
 [tool.poetry.dependencies]
-python = "<3.10,>=3.6.2"
+python = ">=3.6.2"
 requests = "^2.25.1"
 singer-sdk = "^0.3.11"
 boto3 = "^1.18.63"

tap_aws_cost_explorer/schema.py

@@ -0,0 +1,84 @@
+from singer_sdk import typing as th
+
+config_jsonschema = th.PropertiesList(
+    th.Property(
+        "access_key",
+        th.StringType,
+        required=True,
+        description="Your AWS Account Access Key."
+    ),
+    th.Property(
+        "secret_key",
+        th.StringType,
+        required=True,
+        description="Your AWS Account Secret Key."
+    ),
+    th.Property(
+        "session_token",
+        th.StringType,
+        description="Your AWS Account Session Token if required for authentication."
+    ),
+    th.Property(
+        "start_date",
+        th.StringType,
+        required=True,
+        description="The start date for retrieving Amazon Web Services cost."
+    ),
+    th.Property(
+        "end_date",
+        th.DateTimeType,
+        description="The end date for retrieving Amazon Web Services cost."
+    ),
+    th.Property(
+        "granularity",
+        th.StringType,
+        required=True,
+        description="Sets the Amazon Web Services cost granularity to \
+                    MONTHLY or DAILY , or HOURLY."
+    ),
+    th.Property(
+        "metrics",
+        th.ArrayType(th.StringType),
+        required=True,
+        description="Which metrics are returned in the query. Valid \
+                    values are AmortizedCost, BlendedCost, \
+                    NetAmortizedCost, NetUnblendedCost, \
+                    NormalizedUsageAmount, UnblendedCost, and \
+                    UsageQuantity.\
+                    See boto3 docs: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ce/client/get_cost_and_usage_with_resources.html\
+                    "
+    ),
+    th.Property(
+        "stream_name",
+        th.StringType,
+        required = False,
+        description="Name of stream, often results in names of tables objects."
+    ),
+    th.Property(
+        "filter",
+        th.ObjectType(
+            th.Property(
+                "Dimensions",
+                th.ObjectType(
+                    th.Property("Key",th.StringType),
+                    th.Property("Values",th.ArrayType(
+                        th.StringType
+                    ))
+                )
+            )
+        ),
+        required=False,
+        description="See boto3 client explorer params: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ce/client/get_cost_and_usage_with_resources.html"
+    ),
+    th.Property(
+        "groupby",
+        th.ArrayType(
+            th.ObjectType(
+                th.Property("Type",th.StringType),
+                th.Property("Key",th.StringType)
+            )
+        ),
+        required=False,
+        description="See boto3 client explorer params: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ce/client/get_cost_and_usage_with_resources.html"
+    ),
+).to_dict()

tap_aws_cost_explorer/streams.py

@@ -1,6 +1,6 @@
 """Stream type classes for tap-aws-cost-explorer."""
 
-import datetime
+import datetime, json
 from pathlib import Path
 from typing import Optional, Iterable
 
@@ -11,18 +11,24 @@
 
 class CostAndUsageWithResourcesStream(AWSCostExplorerStream):
     """Define custom stream."""
-    name = "cost"
-    primary_keys = ["metric_name", "time_period_start"]
+    primary_keys = ["metric_name", "groupby_values","filter_config","time_period_start"]
     replication_key = "time_period_start"
     # Optionally, you may also use `schema_filepath` in place of `schema`:
     # schema_filepath = SCHEMAS_DIR / "users.json"
     schema = th.PropertiesList(
         th.Property("time_period_start", th.DateTimeType),
         th.Property("time_period_end", th.DateTimeType),
+        th.Property("groupby_keys", th.StringType),
+        th.Property("groupby_values", th.StringType),
+        th.Property("filter_config", th.StringType),
         th.Property("metric_name", th.StringType),
-        th.Property("amount", th.StringType),
         th.Property("amount_unit", th.StringType),
+        th.Property("amount", th.NumberType),
     ).to_dict()
+
+    def __init__(self,tap):
+        self.name = tap.config.get('stream_name','costs')
+        super().__init__(tap)
 
     def _get_end_date(self):
         if self.config.get("end_date") is None:
@@ -34,6 +40,8 @@ def get_records(self, context: Optional[dict]) -> Iterable[dict]:
         next_page = True
         start_date = self.get_starting_timestamp(context)
         end_date = self._get_end_date()
+        filter_config_str = json.dumps(self.config.get('filter',{}))
+        groupby_keys_str = ','.join([i.get('Key') for i in self.config.get('groupby',[]) ])
 
         while next_page:
             response = self.conn.get_cost_and_usage(
@@ -43,15 +51,35 @@ def get_records(self, context: Optional[dict]) -> Iterable[dict]:
                 },
                 Granularity=self.config.get("granularity"),
                 Metrics=self.config.get("metrics"),
+                GroupBy=self.config.get('groupby',[]),
+                Filter=self.config.get('filter',{})
             )
             next_page = response.get("NextPageToken")
 
             for row in response.get("ResultsByTime"):
-                for k, v in row.get("Total").items():
-                    yield {
-                        "time_period_start": row.get("TimePeriod").get("Start"),
-                        "time_period_end": row.get("TimePeriod").get("End"),
-                        "metric_name": k,
-                        "amount": v.get("Amount"),
-                        "amount_unit": v.get("Unit")
-                    }
+                has_groups = len(row.get('Groups',[])) > 0
+                if has_groups:
+                    for group in row['Groups']:
+                        for k,v in group['Metrics'].items():
+                            yield {
+                                "time_period_start": row.get("TimePeriod").get("Start"),
+                                "time_period_end": row.get("TimePeriod").get("End"),
+                                "groupby_keys": groupby_keys_str,
+                                "groupby_values": ','.join(group['Keys']),
+                                "filter_config": filter_config_str,
+                                "metric_name": k,
+                                "amount_unit": v.get("Unit"),
+                                "amount": float(v.get("Amount"))
+                            }
+                else:
+                    for k, v in row.get("Total").items():
+                        yield {
+                            "time_period_start": row.get("TimePeriod").get("Start"),
+                            "time_period_end": row.get("TimePeriod").get("End"),
+                            "metric_name": k,
+                            "groupby_keys": None,
+                            "groupby_values": None,
+                            "filter_config": filter_config_str,
+                            "amount": v.get("Amount"),
+                            "amount_unit": v.get("Unit")
+                        }

tap_aws_cost_explorer/tap.py

@@ -3,68 +3,16 @@
 from typing import List
 
 from singer_sdk import Tap, Stream
-from singer_sdk import typing as th  # JSON schema typing helpers
 
-from tap_aws_cost_explorer.streams import (
-    CostAndUsageWithResourcesStream,
-)
-STREAM_TYPES = [
-    CostAndUsageWithResourcesStream,
-]
+from tap_aws_cost_explorer.schema import config_jsonschema
+from tap_aws_cost_explorer.streams import CostAndUsageWithResourcesStream
 
 
 class TapAWSCostExplorer(Tap):
     """AWSCostExplorer tap class."""
     name = "tap-aws-cost-explorer"
-
-    config_jsonschema = th.PropertiesList(
-        th.Property(
-            "access_key",
-            th.StringType,
-            required=True,
-            description="Your AWS Account Access Key."
-        ),
-        th.Property(
-            "secret_key",
-            th.StringType,
-            required=True,
-            description="Your AWS Account Secret Key."
-        ),
-        th.Property(
-            "session_token",
-            th.StringType,
-            description="Your AWS Account Session Token if required for authentication."
-        ),
-        th.Property(
-            "start_date",
-            th.StringType,
-            required=True,
-            description="The start date for retrieving Amazon Web Services cost."
-        ),
-        th.Property(
-            "end_date",
-            th.DateTimeType,
-            description="The end date for retrieving Amazon Web Services cost."
-        ),
-        th.Property(
-            "granularity",
-            th.StringType,
-            required=True,
-            description="Sets the Amazon Web Services cost granularity to \
-                        MONTHLY or DAILY , or HOURLY."
-        ),
-        th.Property(
-            "metrics",
-            th.ArrayType(th.StringType),
-            required=True,
-            description="Which metrics are returned in the query. Valid \
-                        values are AmortizedCost, BlendedCost, \
-                        NetAmortizedCost, NetUnblendedCost, \
-                        NormalizedUsageAmount, UnblendedCost, and \
-                        UsageQuantity."
-        ),
-    ).to_dict()
+    config_jsonschema = config_jsonschema
 
     def discover_streams(self) -> List[Stream]:
         """Return a list of discovered streams."""
-        return [stream_class(tap=self) for stream_class in STREAM_TYPES]
+        return [CostAndUsageWithResourcesStream(tap=self)]