fix: Add better logging for scheduler

cloudquery/sdk/scheduler/scheduler.py

@@ -10,7 +10,7 @@
     SyncMigrateTableMessage,
 )
 from cloudquery.sdk.schema import Resource
-from .table_resolver import TableResolver
+from .table_resolver import TableResolver, Client
 
 QUEUE_PER_WORKER = 100
 
@@ -72,7 +72,7 @@ def shutdown(self):
             pool.shutdown()
 
     def resolve_resource(
-        self, resolver: TableResolver, client, parent: Resource, item: Any
+        self, resolver: TableResolver, client: Client, parent: Resource, item: Any
     ) -> Resource:
         resource = Resource(resolver.table, parent, item)
         resolver.pre_resource_resolve(client, resource)
@@ -85,19 +85,25 @@ def resolve_table(
         self,
         resolver: TableResolver,
         depth: int,
-        client,
+        client: Client,
         parent_item: Resource,
         res: queue.Queue,
     ):
         table_resolvers_started = 0
         try:
             if depth == 0:
                 self._logger.info(
-                    "table resolver started", table=resolver.table.name, depth=depth
+                    "table resolver started",
+                    client_id=client.id(),
+                    table=resolver.table.name,
+                    depth=depth,
                 )
             else:
                 self._logger.debug(
-                    "table resolver started", table=resolver.table.name, depth=depth
+                    "table resolver started",
+                    client_id=client.id(),
+                    table=resolver.table.name,
+                    depth=depth,
                 )
             total_resources = 0
             for item in resolver.resolve(client, parent_item):
@@ -108,6 +114,7 @@ def resolve_table(
                 except Exception as e:
                     self._logger.error(
                         "failed to resolve resource",
+                        client_id=client.id(),
                         table=resolver.table.name,
                         depth=depth,
                         exc_info=True,
@@ -128,19 +135,25 @@ def resolve_table(
             if depth == 0:
                 self._logger.info(
                     "table resolver finished successfully",
+                    client_id=client.id(),
                     table=resolver.table.name,
+                    resources=total_resources,
                     depth=depth,
                 )
             else:
                 self._logger.debug(
                     "table resolver finished successfully",
+                    client_id=client.id(),
                     table=resolver.table.name,
+                    resources=total_resources,
                     depth=depth,
                 )
         except Exception as e:
             self._logger.error(
                 "table resolver finished with error",
+                client_id=client.id(),
                 table=resolver.table.name,
+                resources=total_resources,
                 depth=depth,
                 exc_info=True,
             )

cloudquery/sdk/scheduler/table_resolver.py

@@ -1,6 +1,11 @@
 from cloudquery.sdk.schema.table import Table
 from cloudquery.sdk.schema import Resource
-from typing import Any, Generator
+from typing import Any, Generator, List
+
+
+class Client:
+    def id(self):
+        raise NotImplementedError()
 
 
 class TableResolver:
@@ -16,22 +21,22 @@ def table(self) -> Table:
     def child_resolvers(self):
         return self._child_resolvers
 
-    def multiplex(self, client):
+    def multiplex(self, client: Client) -> List[Client]:
         return [client]
 
     def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
         raise NotImplementedError()
 
-    def pre_resource_resolve(self, client, resource):
+    def pre_resource_resolve(self, client: Client, resource):
         return
 
-    def resolve_column(self, client, resource: Resource, column_name: str):
+    def resolve_column(self, client: Client, resource: Resource, column_name: str):
         if type(resource.item) is dict:
             if column_name in resource.item:
                 resource.set(column_name, resource.item[column_name])
         else:
             if hasattr(resource.item, column_name):
                 resource.set(column_name, resource.item[column_name])
 
-    def post_resource_resolve(self, client, resource):
+    def post_resource_resolve(self, client: Client, resource):
         return

cloudquery/sdk/serve/plugin.py

@@ -16,7 +16,7 @@
 
 
 def get_logger(args):
-    log = structlog.get_logger()
+    log = structlog.get_logger(processors=[structlog.processors.JSONRenderer()])
     return log
 
 

tests/scheduler/scheduler.py

@@ -37,7 +37,8 @@ def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
 
 
 class TestClient:
-    pass
+    def id(self):
+        return "test_client"
 
 
 def test_scheduler():