wip scheduler

Author: yevgenypats

cloudquery/sdk/scheduler/scheduler.py

@@ -2,10 +2,12 @@
 from typing import List, Generator, Any
 import queue
 import time
-from cloudquery.sdk.message import SyncMessage, SyncInsertMessage
+from cloudquery.sdk.schema import Table
+from cloudquery.sdk.message import SyncMessage, SyncInsertMessage, SyncMigrateMessage
 import concurrent.futures
 from typing import Generator
 
+# This is all WIP
 class Task:
    def __init__(self, fetcher, parent_item):
      self._fetcher = fetcher
@@ -66,15 +68,47 @@ def run_task(task):
 
     print(f"Running task {task}")
 
+
+QUEUE_PER_WORKER = 100
+
 class Scheduler:
-    def __init__(self):
+    def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3):
         self._queue = queue.Queue()
-        # self._pool = ThreadPool(processes=10)
+        self._max_depth = max_depth
+        if concurrency <= 0:
+           raise ValueError("concurrency must be greater than 0")
+        if queue_size <= 0:
+            raise ValueError("queue_size must be greater than 0")
+        if max_depth <= 0:
+            raise ValueError("max_depth must be greater than 0")
+        self._queue_size = queue_size if queue_size > 0 else concurrency * QUEUE_PER_WORKER
+        self._pools = []
+        self._queues = []
+        current_depth_concurrency = concurrency
+        current_depth_queue_size = queue_size
+        for i in range(max_depth + 1):
+          self._queues.append(queue.Queue(maxsize=current_depth_queue_size))
+          self._pools.append(concurrent.futures.ThreadPoolExecutor(max_workers=current_depth_concurrency))
+          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 worker(self, max_depth: int):
+      while True:
+        task = self._queues[max_depth].get()
+        if task is None:
+            break
+        self._pools[max_depth].submit(*task)
+    
+    def table_resolver(self, table: Table, client, res: queue.Queue):
+      for resource in table.resolve(client):
+         pass
+      #  task.resolve
 
-    def sync(self, fetchers: List[Fetcher]) -> Generator[SyncMessage]:
-      task_queue = queue.Queue()
-      with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-          for worker_id in range(3):
-              executor.submit(worker_task, task_queue, worker_id)
-      main_task(task_queue)
-        
+    def sync(self, client, tables: List[Table], res: queue.Queue, deterministic_cq_id=False):
+      for table in tables:
+        res.put(SyncMigrateMessage(record=table.to_arrow_schemas()))
+      for table in tables:
+         clients = table.multiplex(client)
+         for client in clients:
+          self._queues[0].put((table.resolver, client, res))
+      self._queues[0].put(None)

cloudquery/sdk/schema/table.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-from typing import List
+from typing import List, Generator, Any
 
 import pyarrow as pa
 
 from cloudquery.sdk.schema import arrow
 from .column import Column
 
+class Client:
+    pass
 
 class Table:
     def __init__(self, name: str, columns: List[Column], title: str = "", description: str = "",
@@ -21,6 +23,12 @@ def __init__(self, name: str, columns: List[Column], title: str = "", descriptio
         self.relations = relations
         self.is_incremental = is_incremental
 
+    def multiplex(self, client) -> List[Table]:
+      raise [client]
+
+    def resolver(self, client: Client, parent=None) -> Generator[Any]:
+      raise NotImplementedError()
+
     @property
     def primary_keys(self):
         return [column.name for column in self.columns if column.primary_key]