wip scheduler

Author: yevgenypats

cloudquery/sdk/scalar/__init__.py

@@ -0,0 +1,6 @@
+from .scalar import Scalar, ScalarInvalidTypeError, ScalarFactory
+from .binary import Binary
+from .bool import Bool
+from .date32 import Date32
+from .float64 import Float64
+from .int64 import Int64

cloudquery/sdk/scalar/binary.py

@@ -0,0 +1,26 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+class Binary(Scalar):
+    def __init__(self, valid: bool = False, value: bytes = None):
+      self._valid = valid
+      self._value = value 
+
+    def __eq__(self, scalar: Scalar) -> bool:
+      if scalar is None:
+          return False
+      if type(scalar) == Binary:
+          return self._value == scalar._value and self._valid == scalar._valid
+      return False
+    
+    def set(self, scalar):
+      if scalar is None:
+          return
+
+      if type(scalar) == bytes:
+          self._valid = True
+          self._value = scalar
+      elif type(scalar) == str:
+          self._valid = True
+          self._value = scalar.encode()
+      else:
+          raise ScalarInvalidTypeError("Invalid type for Binary scalar")

cloudquery/sdk/scalar/bool.py

@@ -0,0 +1,41 @@
+
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from typing import Any
+
+def parse_string_to_bool(input_string):
+    true_strings = ['true', 't', 'yes', 'y', '1']
+    false_strings = ['false', 'f', 'no', 'n', '0']
+
+    lower_input = input_string.lower()
+
+    if lower_input in true_strings:
+        return True
+    elif lower_input in false_strings:
+        return False
+    else:
+        raise ScalarInvalidTypeError("Invalid boolean string: {}".format(input_string))
+
+class Bool(Scalar):
+    def __init__(self, valid: bool = False, value: bool = False) -> None:
+      self._valid = valid
+      self._value = value
+    
+    def __eq__(self, scalar: Scalar) -> bool:
+      if scalar is None:
+          return False
+      if type(scalar) == Bool:
+          return self._value == scalar._value and self._valid == scalar._valid
+      return False
+    
+    def set(self, value: Any):
+      if value is None:
+          return
+
+      if type(value) == bool:
+        self._value = value
+      elif type(value) == str:
+        self._value = parse_string_to_bool(value)
+      else:
+        raise ScalarInvalidTypeError("Invalid type for Bool scalar")
+      
+      self._valid = True

cloudquery/sdk/scalar/date32.py

@@ -0,0 +1,31 @@
+
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from datetime import datetime, time
+from typing import Any
+
+class Date32(Scalar):
+    def __init__(self, valid: bool = False, value: bool = False) -> None:
+      self._valid = valid
+      self._value = value
+    
+    def __eq__(self, scalar: Scalar) -> bool:
+      if scalar is None:
+          return False
+      if type(scalar) == Date32:
+          return self._value == scalar._value and self._valid == scalar._valid
+      return False
+    
+    def set(self, value: Any):
+      if value is None:
+          return
+
+      if type(value) == datetime:
+        self._value = value
+      elif type(value) == str:
+        self._value = datetime.strptime(value, "%Y-%m-%d")
+      elif type(value) == time:
+        self._value = datetime.combine(datetime.today(), value)
+      else:
+        raise ScalarInvalidTypeError("Invalid type for Bool scalar")
+      
+      self._valid = True

cloudquery/sdk/scalar/float64.py

@@ -0,0 +1,30 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+class Float64(Scalar):
+    def __init__(self, valid: bool = False, value: float = None):
+      self._valid = valid
+      self._value = value 
+
+    def __eq__(self, scalar: Scalar) -> bool:
+      if scalar is None:
+          return False
+      if type(scalar) == Float64:
+          return self._value == scalar._value and self._valid == scalar._valid
+      return False
+    
+    def set(self, value):
+      if value is None:
+          return
+
+      if type(value) == int:
+          self._value = float(value)
+      elif type(value) == float:
+          self._value = value
+      elif type(value) == str:
+          try:
+            self._value = float(value)
+          except ValueError:
+            raise ScalarInvalidTypeError("Invalid type for Float64 scalar")
+      else:
+          raise ScalarInvalidTypeError("Invalid type for Binary scalar")
+      self._valid = True

cloudquery/sdk/scalar/int64.py

@@ -0,0 +1,30 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+class Int64(Scalar):
+    def __init__(self, valid: bool = False, value: float = None):
+      self._valid = valid
+      self._value = value 
+
+    def __eq__(self, scalar: Scalar) -> bool:
+      if scalar is None:
+          return False
+      if type(scalar) == Int64:
+          return self._value == scalar._value and self._valid == scalar._valid
+      return False
+    
+    def set(self, value):
+      if value is None:
+          return
+
+      if type(value) == int:
+          self._value = value
+      elif type(value) == float:
+          self._value = int(value)
+      elif type(value) == str:
+          try:
+            self._value = int(value)
+          except ValueError:
+            raise ScalarInvalidTypeError("Invalid type for Int64 scalar")
+      else:
+          raise ScalarInvalidTypeError("Invalid type for Int64 scalar")
+      self._valid = True

cloudquery/sdk/scalar/scalar.py

@@ -0,0 +1,23 @@
+import pyarrow as pa
+from .int64 import Int64
+
+class ScalarInvalidTypeError(Exception):
+  pass
+
+class Scalar:
+    @property
+    def is_valid(self) -> bool:
+      return self._valid
+
+
+class ScalarFactory:
+  def __init__(self):
+    self._type_map = {
+        pa.int64: lambda dt: Int64(),
+    }
+  
+  def new_scalar(self, dt):
+    if dt in self._type_map:
+      return self._type_map[dt]()
+    else:
+      raise ScalarInvalidTypeError("Invalid type for scalar")

cloudquery/sdk/scheduler/scheduler.py

@@ -2,75 +2,19 @@
 from typing import List, Generator, Any
 import queue
 import time
-from cloudquery.sdk.schema import Table
+from cloudquery.sdk.schema import Table, Resource
 from cloudquery.sdk.message import SyncMessage, SyncInsertMessage, SyncMigrateMessage
-import concurrent.futures
+from concurrent import futures
 from typing import Generator
-
-# This is all WIP
-class Task:
-   def __init__(self, fetcher, parent_item):
-     self._fetcher = fetcher
-     self._parent_item = parent_item
-
-class Item:
-   def __init__(self, fetcher, parent_item):
-     self._fetcher = fetcher
-     self._parent_item = parent_item
-
-class Fetcher:
-    def __init__(self, relations: List[Any]):
-        self._relations = relations
-    
-    def get(self, parent_item) -> Generator[SyncInsertMessage]:
-        for i in range(10):
-           yield SyncInsertMessage(None)
-    
-    def process_item(self, item: Any) -> Generator[SyncInsertMessage]:
-       pass
-    
-    def transform_item(self, item: Any) -> Any:
-       pass
-        
-
-def worker(fetcher: Fetcher, parent_item: Any):
-  for arr in fetcher.get(parent_item):
-     for res in arr:
-        fetcher.get()
-        
-  # while True:
-  #   item = q.get()
-  #   if item is None:
-  #       break
-  #   do_work(item)
-  #   q.task_done()
-
-def worker_task(q, worker_id):
-    for i in range(5):
-        time.sleep(0.1)  # Simulate work
-        task = (worker_id, i)
-        print(f"Worker {worker_id} created task {task}")
-        q.put(task)
-
-def main_task(q: queue.Queue):
-    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-        while True:
-            try:
-                task = q.get(timeout=1)  # Wait up to 1 second for a task.
-            except queue.Empty:
-                print("All tasks completed.")
-                break  # Exit while loop if no more tasks.
-            print(f"Main: Running task {task}")
-            executor.submit(run_task, task)
-            q.task_done()
-
-def run_task(task):
-    
-    print(f"Running task {task}")
-
+from .table_resolver import TableResolver
 
 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 Scheduler:
     def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3):
         self._queue = queue.Queue()
@@ -82,33 +26,50 @@ def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3):
         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 = []
+        self._pools : List[ThreadPoolExecutorWithQueueSizeLimit] = []
         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))
+        for _ in range(max_depth + 1):
+          self._pools.append(ThreadPoolExecutorWithQueueSizeLimit(maxsize=current_depth_queue_size,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 resolve_resource(self, resolver: TableResolver, client, parent: Resource, item: Any) -> Resource:
+      resource = Resource(resolver.table, None, 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):
+      for item in resolver.resolve(client, parent_item):
+        resource = self.resolve_resource(resolver, client, parent_item)
+        res.put(SyncInsertMessage(resource))
+      res.put(None)
 
-    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)
+    def _sync(self, client, resolvers: List[TableResolver], res: queue.Queue, deterministic_cq_id=False):
+      internal_res = queue.Queue()
+      for resolver in resolvers:
+        clients = resolver.multiplex(client)
+        for client in clients:
+          self._pools[0].submit(self.resolve_table, resolver, client, None, internal_res)
+      while True:
+        message = internal_res.get()
+        if message is None:
+          break
+        res.put(message)
+      res.put(None)
+       
+    def sync(self, client, resolvers: List[TableResolver], deterministic_cq_id=False) -> Generator[SyncMessage]:
+      res = queue.Queue()
+      for resolver in resolvers:
+        yield SyncMigrateMessage(record=resolver.table.to_arrow_schemas())
+      thread = futures.ThreadPoolExecutor()
+      thread.submit(self._sync, client, resolvers, res, deterministic_cq_id)
+      while True:
+        message = res.get()
+        if message is None:
+          break
+        yield message
+      thread.shutdown()

cloudquery/sdk/scheduler/table_resolver.py

@@ -0,0 +1,28 @@
+
+from cloudquery.sdk.schema.table import Table
+from cloudquery.sdk.schema import Resource
+from typing import Any
+
+class TableResolver:
+    def __init__(self, table: Table) -> None:
+        self._table = table
+    
+    @property
+    def table(self):
+      return self._table
+
+    def multiplex(self, client):
+      return [client]
+    
+    def resolve(self, client, parent_resource) -> Any:
+      raise NotImplementedError()
+
+    def pre_resource_resolve(self, client, resource):
+      return
+    
+    def resolve_column(self, client, resource: Resource, column_name: str):
+      if hasattr(resource.item, column_name):
+         resource.set(column_name, resource.item[column_name])
+    
+    def post_resource_resolve(self, client, resource):
+      return

cloudquery/sdk/schema/__init__.py

@@ -1,2 +1,4 @@
 from .column import Column
 from .table import Table, tables_to_arrow_schemas
+from .resource import Resource
+# from .table_resolver import TableReso