Merge pull request #6 from SpareCores/DEV-21

DEV_21

README.md

@@ -92,36 +92,3 @@ server = session.exec(select(Server).where(Server.id == 'trn1.32xlarge')).one()
 pp(server)
 pp(server.vendor)
 ```
-
-Lower level access examples:
-
-```py
-from sc_crawler.vendors import aws
-
-# enable persistent caching of AWS queries
-from cachier import set_default_params
-set_default_params(caching_enabled=True)
-
-# fetch data
-aws.get_all()  # slow to query all instance types in all regions
-
-# look around
-aws.datacenters
-aws.zones
-
-# pretty printed objects
-from rich import print as pp
-pp(aws)
-pp(aws.datacenters)
-pp(aws.servers[0])
-```
-
-Debug raw AWS responses:
-
-```py
-products = aws._methods.get_products()
-pp(products[1]["product"])
-
-instance_types = aws._methods.describe_instance_types(region="us-west-2")
-pp(instance_types[1])
-```

src/sc_crawler/cli.py

@@ -11,6 +11,7 @@
 
 from . import vendors as vendors_module
 from .logger import logger
+from .lookup import compliance_frameworks, countries
 from .schemas import Vendor
 from .utils import hash_database
 
@@ -36,6 +37,9 @@
 log_levels = list(logging._nameToLevel.keys())
 LogLevels = Enum("LOGLEVELS", {k: k for k in log_levels})
 
+supported_tables = [m[10:] for m in dir(Vendor) if m.startswith("inventory_")]
+Tables = Enum("TABLES", {k: k for k in supported_tables})
+
 
 @cli.command()
 def schema(dialect: Engines):
@@ -76,6 +80,10 @@ def pull(
         List[Vendors],
         typer.Option(help="Exclude specific vendor. Can be specified multiple times."),
     ] = [],
+    update_table: Annotated[
+        List[Tables],
+        typer.Option(help="Tables to be updated. Can be specified multiple times."),
+    ] = supported_tables,
     log_level: Annotated[
         LogLevels, typer.Option(help="Log level threshold.")
     ] = LogLevels.INFO.value,  # TODO drop .value after updating Enum to StrEnum in Python3.11
@@ -108,7 +116,7 @@ def custom_serializer(x):
     # enable logging
     channel = logging.StreamHandler()
     formatter = logging.Formatter(
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        "%(asctime)s - %(name)s/%(module)s:%(funcName)s - %(levelname)s - %(message)s"
     )
     channel.setFormatter(formatter)
     logger.setLevel(log_level.value)
@@ -128,10 +136,35 @@ def custom_serializer(x):
     engine = create_engine(connection_string, json_serializer=custom_serializer)
     SQLModel.metadata.create_all(engine)
     with Session(engine) as session:
+        # add/merge static objects to database
+        for compliance_framework in compliance_frameworks.values():
+            session.merge(compliance_framework)
+        for country in countries.values():
+            session.merge(country)
+        # get data for each vendor and then add/merge to database
         for vendor in vendors:
             logger.info("Starting to collect data from vendor: " + vendor.id)
-            vendor.get_all()
-            session.add(vendor)
+            vendor = session.merge(vendor)
+            vendor.set_session(session)
+            if Tables.compliance_frameworks in update_table:
+                vendor.inventory_compliance_frameworks()
+            if Tables.datacenters in update_table:
+                vendor.inventory_datacenters()
+            if Tables.zones in update_table:
+                vendor.inventory_zones()
+            if Tables.servers in update_table:
+                vendor.inventory_servers()
+            if Tables.server_prices in update_table:
+                vendor.inventory_server_prices()
+            if Tables.server_prices_spot in update_table:
+                vendor.inventory_server_prices_spot()
+            if Tables.storage_prices in update_table:
+                vendor.inventory_storage_prices()
+            if Tables.traffic_prices in update_table:
+                vendor.inventory_traffic_prices()
+            if Tables.ipv4_prices in update_table:
+                vendor.inventory_ipv4_prices()
+            session.merge(vendor)
             session.commit()
 
 

src/sc_crawler/logger.py

@@ -2,3 +2,20 @@
 
 logger = logging.getLogger("sc_crawler")
 logger.addHandler(logging.NullHandler())
+
+
+def log_start_end(func):
+    """Log the start and end of the decorated function."""
+
+    def wrap(*args, **kwargs):
+        try:
+            self = args[0]
+            fname = f"{self.id}/{func.__name__}"
+        except Exception:
+            fname = func.__name__
+        logger.debug("Starting %s", fname)
+        result = func(*args, **kwargs)
+        logger.debug("Finished %s", fname)
+        return result
+
+    return wrap