some more work on scheduler with relations

Author: yevgenypats

cloudquery/sdk/internal/servers/plugin_v3/plugin.py

@@ -1,4 +1,5 @@
 import pyarrow as pa
+import structlog
 
 from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc
 from cloudquery.sdk.message import SyncInsertMessage, SyncMigrateTableMessage
@@ -7,7 +8,8 @@
 
 
 class PluginServicer(plugin_pb2_grpc.PluginServicer):
-    def __init__(self, plugin: Plugin):
+    def __init__(self, plugin: Plugin, logger=None):
+        self._logger = logger if logger is not None else structlog.get_logger()
         self._plugin = plugin
 
     def GetName(self, request, context):

cloudquery/sdk/scheduler/scheduler.py

@@ -2,6 +2,7 @@
 from typing import List, Generator, Any
 import queue
 import time
+import structlog
 from enum import Enum
 from cloudquery.sdk.schema import Table, Resource
 from cloudquery.sdk.message import SyncMessage, SyncInsertMessage, SyncMigrateTableMessage
@@ -12,30 +13,33 @@
 
 QUEUE_PER_WORKER = 100
 
+
 class ThreadPoolExecutorWithQueueSizeLimit(futures.ThreadPoolExecutor):
     def __init__(self, maxsize, *args, **kwargs):
         super(ThreadPoolExecutorWithQueueSizeLimit, self).__init__(*args, **kwargs)
         self._work_queue = queue.Queue(maxsize=maxsize)
 
 
-class WorkerStatus:
-  def __init__(self, total_table_resolvers) -> None:
-    self._total_table_resolvers = total_table_resolvers
+class TableResolverStarted:
+  def __init__(self, count=1) -> None:
+    self._count = count
 
   @property
-  def total_table_resolvers(self):
-   return self._total_table_resolvers
+  def count(self):
+   return self._count
 
 
-class TableResolverStatus:
+class TableResolverFinished:
   def __init__(self) -> None:
       pass
 
 
 class Scheduler:
-    def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3):
+    def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3, logger=None):
         self._queue = queue.Queue()
         self._max_depth = max_depth
+        if logger is None:
+          self._logger = structlog.get_logger()
         if concurrency <= 0:
            raise ValueError("concurrency must be greater than 0")
         if max_depth <= 0:
@@ -49,34 +53,53 @@ def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3):
           current_depth_concurrency = current_depth_concurrency // 2 if current_depth_concurrency > 1 else 1
           current_depth_queue_size = current_depth_queue_size // 2 if current_depth_queue_size > 1 else 1
 
+    def shutdown(self):
+      for pool in self._pools:
+        pool.shutdown()
+
     def resolve_resource(self, resolver: TableResolver, client, parent: Resource, item: Any) -> Resource:
-      resource = Resource(resolver.table, None, item)
+      resource = Resource(resolver.table, parent, item)
       resolver.pre_resource_resolve(client, resource)
       for column in resolver.table.columns:
         resolver.resolve_column(client, resource, column.name)
       resolver.post_resource_resolve(client, resource)
       return resource
 
-    def resolve_table(self, resolver: TableResolver, client, parent_item: Any, res: queue.Queue):
+    def resolve_table(self, resolver: TableResolver, depth: int, 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)
+        else:
+          self._logger.debug("table resolver started", table=resolver.table.name, depth=depth)
+        total_resources = 0
         for item in resolver.resolve(client, parent_item):
           resource = self.resolve_resource(resolver, client, parent_item, item)
           res.put(SyncInsertMessage(resource.to_arrow_record()))
+          for child_resolvers in resolver.child_resolvers:
+            self._pools[depth + 1].submit(self.resolve_table, child_resolvers, depth + 1, client, resource, res)
+            table_resolvers_started += 1
+          total_resources += 1
+        if depth == 0:
+          self._logger.info("table resolver finished successfully", table=resolver.table.name, depth=depth)
+        else:
+          self._logger.debug("table resolver finished successfully", table=resolver.table.name, depth=depth)
       except Exception as e:
-        traceback.print_exc()
-        print("exception")
-        print(e)
+        self._logger.error("table resolver finished with error", table=resolver.table.name, depth=depth, exception=e)
       finally:
-        res.put(TableResolverStatus())
+        res.put(TableResolverStarted(count=table_resolvers_started))
+        res.put(TableResolverFinished())
 
     def _sync(self, client, resolvers: List[TableResolver], res: queue.Queue, deterministic_cq_id=False):
       total_table_resolvers = 0
-      for resolver in resolvers:
-        clients = resolver.multiplex(client)
-        for client in clients:
-          self._pools[0].submit(self.resolve_table, resolver, client, None, res)
-          total_table_resolvers += 1
-      res.put(WorkerStatus(total_table_resolvers))
+      try:
+        for resolver in resolvers:
+          clients = resolver.multiplex(client)
+          for client in clients:
+            self._pools[0].submit(self.resolve_table, resolver, 0, client, None, res)
+            total_table_resolvers += 1
+      finally:
+        res.put(TableResolverStarted(total_table_resolvers))
 
     def sync(self, client, resolvers: List[TableResolver], deterministic_cq_id=False) -> Generator[SyncMessage, None, None]:
       res = queue.Queue()
@@ -88,12 +111,12 @@ def sync(self, client, resolvers: List[TableResolver], deterministic_cq_id=False
       finished_table_resovlers = 0
       while True:
         message = res.get()
-        if type(message) == WorkerStatus:
-          total_table_resolvers += message.total_table_resolvers
+        if type(message) == TableResolverStarted:
+          total_table_resolvers += message.count
           if total_table_resolvers == finished_table_resovlers:
             break
           continue
-        elif type(message) == TableResolverStatus:
+        elif type(message) == TableResolverFinished:
           finished_table_resovlers += 1
           if total_table_resolvers == finished_table_resovlers:       
             break

cloudquery/sdk/scheduler/table_resolver.py

@@ -4,13 +4,18 @@
 from typing import Any,Generator
 
 class TableResolver:
-    def __init__(self, table: Table) -> None:
+    def __init__(self, table: Table, child_resolvers=[]) -> None:
         self._table = table
+        self._child_resolvers = child_resolvers
 
     @property
     def table(self) -> Table:
       return self._table
 
+    @property
+    def child_resolvers(self):
+      return self._child_resolvers
+
     def multiplex(self, client):
       return [client]
 
@@ -21,7 +26,11 @@ def pre_resource_resolve(self, client, resource):
       return
 
     def resolve_column(self, client, resource: Resource, column_name: str):
-      if hasattr(resource.item, column_name):
+      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):

cloudquery/sdk/serve/plugin.py

@@ -1,4 +1,5 @@
 import argparse
+import structlog
 from concurrent import futures
 
 import grpc
@@ -14,6 +15,10 @@
 DOC_FORMATS = ["json", "markdown"]
 
 
+def get_logger(args):
+    log = structlog.get_logger()
+    return log
+
 class PluginCommand:
     def __init__(self, plugin: Plugin):
         self._plugin = plugin
@@ -57,11 +62,12 @@ def run(self, args):
             sys.exit(1)
 
     def _serve(self, args):
+        logger = get_logger(args)
         server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
         discovery_pb2_grpc.add_DiscoveryServicer_to_server(
             DiscoveryServicer([3]), server)
         plugin_pb2_grpc.add_PluginServicer_to_server(
-            PluginServicer(self._plugin), server)
+            PluginServicer(self._plugin, logger), server)
         server.add_insecure_port(args.address)
         print("Starting server. Listening on " + args.address)
         server.start()

tests/scheduler/scheduler.py

@@ -1,6 +1,7 @@
 
 from typing import Any, List, Generator
 import pyarrow as pa
+import pytest
 from cloudquery.sdk.scheduler import Scheduler, TableResolver
 from cloudquery.sdk.schema import Table, Column, Resource
 from cloudquery.sdk.message import SyncMessage
@@ -12,22 +13,40 @@ def __init__(self):
             Column("test_column", pa.int64())
         ])
 
+class SchedulerTestChildTable(Table):
+    def __init__(self):
+        super().__init__("test_child_table", [
+            Column("test_child_column", pa.int64())
+        ])
+
 class SchedulerTestTableResolver(TableResolver):
     def __init__(self) -> None:
-        super().__init__(SchedulerTestTable())
+        super().__init__(SchedulerTestTable(), child_resolvers=[SchedulerTestChildTableResolver()])
 
     def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
       yield {"test_column": 1}
 
+class SchedulerTestChildTableResolver(TableResolver):
+    def __init__(self) -> None:
+        super().__init__(SchedulerTestChildTable())
+    
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+      yield {"test_child_column": 2}
+
 class TestClient:
     pass
 
 def test_scheduler():
     client = TestClient()
     s = Scheduler(10)
+    table1 = Table("test_table", [Column("test_column", pa.int64())])
+    expected_record1 = pa.record_batch([[1]], schema=table1.to_arrow_schema())
+    table2 = Table("test_child_table", [Column("test_child_column", pa.int64())])
+    expected_record2 = pa.record_batch([[2]], schema=table2.to_arrow_schema())
     resources: List[SyncMessage] = []
     for resource in s.sync(client, [SchedulerTestTableResolver()]):
         resources.append(resource)
-    assert len(resources) == 2
-    print(resources[1].record.to_pydict())
-    
+    assert len(resources) == 3
+    assert resources[1].record == expected_record1
+    assert resources[2].record == expected_record2
+    s.shutdown()