traffic price for both directions

SpareCores · Feb 27, 2024 · afb5a4c · afb5a4c
1 parent dd35080
commit afb5a4c
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 31 deletions.
diff --git a/src/sc_crawler/schemas.py b/src/sc_crawler/schemas.py
@@ -827,7 +827,8 @@ class StoragePrice(StoragePriceBase, table=True):
 
 class TrafficPriceBase(HasDatacenterPK, HasVendorPK):
     direction: TrafficDirection = Field(
-        description="Direction of the traffic: inbound or outbound."
+        description="Direction of the traffic: inbound or outbound.",
+        primary_key=True,
     )
     status: Status = Field(
         default=Status.ACTIVE,

diff --git a/src/sc_crawler/vendors/aws.py b/src/sc_crawler/vendors/aws.py
@@ -792,38 +792,39 @@ def inventory_storage_prices(vendor):
 
 
 def inventory_traffic_prices(vendor):
-    vendor.progress_tracker.start_task(name="Searching for Traffic prices", n=None)
-
     loc2dc = _location_datacenter_map(vendor)
-    products = _boto_get_products(
-        service_code="AWSDataTransfer",
-        filters={
-            "transferType": "AWS Inbound",
-        },
-    )
-
-    vendor.progress_tracker.update_task(
-        description="Syncing inbound Traffic prices", total=len(products)
-    )
-    for product in products:
-        try:
-            datacenter = loc2dc[product["product"]["attributes"]["toLocation"]]
-        except KeyError:
-            continue
-        finally:
-            vendor.progress_tracker.advance_task()
-        price = _extract_ondemand_prices(product["terms"])
-        TrafficPrice(
-            vendor=vendor,
-            datacenter=datacenter,
-            price=price[0][-1].get("price"),
-            price_tiered=price,
-            currency=price[1],
-            unit="GB",
-            direction=TrafficDirection.IN,
+    for direction in list(TrafficDirection):
+        vendor.progress_tracker.start_task(
+            name=f"Searching for {direction.value} Traffic prices", n=None
         )
-    vendor.progress_tracker.hide_task()
-    vendor.log(f"{len(products)} IPv4 prices synced.")
+        products = _boto_get_products(
+            service_code="AWSDataTransfer",
+            filters={
+                "transferType": "AWS " + direction.value.title(),
+            },
+        )
+        vendor.progress_tracker.update_task(
+            description=f"Syncing {direction.value} Traffic prices", total=len(products)
+        )
+        for product in products:
+            try:
+                datacenter = loc2dc[product["product"]["attributes"]["toLocation"]]
+            except KeyError:
+                continue
+            finally:
+                vendor.progress_tracker.advance_task()
+            price = _extract_ondemand_prices(product["terms"])
+            TrafficPrice(
+                vendor=vendor,
+                datacenter=datacenter,
+                price=price[0][-1].get("price"),
+                price_tiered=price,
+                currency=price[1],
+                unit="GB",
+                direction=direction,
+            )
+        vendor.progress_tracker.hide_task()
+        vendor.log(f"{len(products)} {direction.value} Traffic prices synced.")
 
 
 def inventory_ipv4_prices(vendor):