working scheduler

yevgenypats · yevgenypats · commit ac5089c4353d · 2023-07-26T14:36:19.000+03:00
diff --git a/cloudquery/sdk/message/sync.py b/cloudquery/sdk/message/sync.py
@@ -11,5 +11,5 @@ def __init__(self, record: pa.RecordBatch):
 
 
 class SyncMigrateTableMessage:
-    def __init__(self, table: pa.Schema):
-        self.table = table
+    def __init__(self, schema: pa.Schema):
+        self.schema = schema
diff --git a/cloudquery/sdk/scalar/__init__.py b/cloudquery/sdk/scalar/__init__.py
@@ -1,4 +1,5 @@
-from .scalar import Scalar, ScalarInvalidTypeError, ScalarFactory
+from .scalar import Scalar, ScalarInvalidTypeError
+from .scalar_factory import ScalarFactory
 from .binary import Binary
 from .bool import Bool
 from .date32 import Date32
diff --git a/cloudquery/sdk/scalar/binary.py b/cloudquery/sdk/scalar/binary.py
@@ -11,6 +11,10 @@ def __eq__(self, scalar: Scalar) -> bool:
       if type(scalar) == Binary:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
+
+    @property
+    def value(self):
+       return self._value
     
     def set(self, scalar):
       if scalar is None:
diff --git a/cloudquery/sdk/scalar/bool.py b/cloudquery/sdk/scalar/bool.py
@@ -26,7 +26,11 @@ def __eq__(self, scalar: Scalar) -> bool:
       if type(scalar) == Bool:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
-    
+
+    @property
+    def value(self):
+       return self._value
+     
     def set(self, value: Any):
       if value is None:
           return
diff --git a/cloudquery/sdk/scalar/date32.py b/cloudquery/sdk/scalar/date32.py
@@ -14,6 +14,10 @@ def __eq__(self, scalar: Scalar) -> bool:
       if type(scalar) == Date32:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
+
+    @property
+    def value(self):
+       return self._value
     
     def set(self, value: Any):
       if value is None:
diff --git a/cloudquery/sdk/scalar/float64.py b/cloudquery/sdk/scalar/float64.py
@@ -11,6 +11,10 @@ def __eq__(self, scalar: Scalar) -> bool:
       if type(scalar) == Float64:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
+
+    @property
+    def value(self):
+       return self._value
     
     def set(self, value):
       if value is None:
diff --git a/cloudquery/sdk/scalar/int64.py b/cloudquery/sdk/scalar/int64.py
@@ -11,7 +11,11 @@ def __eq__(self, scalar: Scalar) -> bool:
       if type(scalar) == Int64:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
-    
+
+    @property
+    def value(self):
+       return self._value
+
     def set(self, value):
       if value is None:
           return
@@ -23,8 +27,8 @@ def set(self, value):
       elif type(value) == str:
           try:
             self._value = int(value)
-          except ValueError:
-            raise ScalarInvalidTypeError("Invalid type for Int64 scalar")
+          except ValueError as e:
+            raise ScalarInvalidTypeError("Invalid type for Int64 scalar") from e
       else:
-          raise ScalarInvalidTypeError("Invalid type for Int64 scalar")
+          raise ScalarInvalidTypeError("Invalid type {} for Int64 scalar".format(type(value)))
       self._valid = True
diff --git a/cloudquery/sdk/scalar/scalar.py b/cloudquery/sdk/scalar/scalar.py
@@ -1,5 +1,3 @@
-import pyarrow as pa
-from .int64 import Int64
 
 class ScalarInvalidTypeError(Exception):
   pass
@@ -8,16 +6,9 @@ class Scalar:
     @property
     def is_valid(self) -> bool:
       return self._valid
+    
+    @property
+    def value(self):
+       raise NotImplementedError("Scalar value not implemented")
 
 
-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")
diff --git a/cloudquery/sdk/scalar/scalar_factory.py b/cloudquery/sdk/scalar/scalar_factory.py
@@ -0,0 +1,15 @@
+
+import pyarrow as pa
+from .scalar import ScalarInvalidTypeError
+from .int64 import Int64
+
+class ScalarFactory:
+  def __init__(self):
+    pass
+
+  def new_scalar(self, dt):
+    dt_id = dt.id
+    if dt_id == pa.types.lib.Type_INT64:
+      return Int64()
+    else:
+      raise ScalarInvalidTypeError("Invalid type {} for scalar".format(dt))
diff --git a/cloudquery/sdk/scheduler/__init__.py b/cloudquery/sdk/scheduler/__init__.py
@@ -0,0 +1,2 @@
+from .scheduler import Scheduler
+from .table_resolver import TableResolver
diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py
@@ -2,11 +2,13 @@
 from typing import List, Generator, Any
 import queue
 import time
+from enum import Enum
 from cloudquery.sdk.schema import Table, Resource
-from cloudquery.sdk.message import SyncMessage, SyncInsertMessage, SyncMigrateMessage
+from cloudquery.sdk.message import SyncMessage, SyncInsertMessage, SyncMigrateTableMessage
 from concurrent import futures
 from typing import Generator
 from .table_resolver import TableResolver
+import traceback
 
 QUEUE_PER_WORKER = 100
 
@@ -15,14 +17,27 @@ 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
+  
+  @property
+  def total_table_resolvers(self):
+   return self._total_table_resolvers
+
+
+class TableResolverStatus:
+  def __init__(self) -> None:
+      pass
+
+
 class Scheduler:
     def __init__(self, concurrency: int, queue_size: int = 0, max_depth : int = 3):
         self._queue = queue.Queue()
         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
@@ -43,33 +58,45 @@ def resolve_resource(self, resolver: TableResolver, client, parent: Resource, it
       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)
+      try:
+        for item in resolver.resolve(client, parent_item):
+          resource = self.resolve_resource(resolver, client, parent_item, item)
+          res.put(SyncInsertMessage(resource.to_arrow_record()))
+      except Exception as e:
+        traceback.print_exc()
+        print("exception")
+        print(e)
+      finally:
+        res.put(TableResolverStatus())
     
     def _sync(self, client, resolvers: List[TableResolver], res: queue.Queue, deterministic_cq_id=False):
-      internal_res = queue.Queue()
+      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, internal_res)
-      while True:
-        message = internal_res.get()
-        if message is None:
-          break
-        res.put(message)
-      res.put(None)
+          self._pools[0].submit(self.resolve_table, resolver, client, None, res)
+          total_table_resolvers += 1
+      res.put(WorkerStatus(total_table_resolvers))
        
-    def sync(self, client, resolvers: List[TableResolver], deterministic_cq_id=False) -> Generator[SyncMessage]:
+    def sync(self, client, resolvers: List[TableResolver], deterministic_cq_id=False) -> Generator[SyncMessage, None, None]:
       res = queue.Queue()
       for resolver in resolvers:
-        yield SyncMigrateMessage(record=resolver.table.to_arrow_schemas())
+        yield SyncMigrateTableMessage(schema=resolver.table.to_arrow_schema())
       thread = futures.ThreadPoolExecutor()
       thread.submit(self._sync, client, resolvers, res, deterministic_cq_id)
+      total_table_resolvers = 0
+      finished_table_resovlers = 0
       while True:
         message = res.get()
-        if message is None:
-          break
+        if type(message) == WorkerStatus:
+          total_table_resolvers += message.total_table_resolvers
+          if total_table_resolvers == finished_table_resovlers:
+            break
+          continue
+        elif type(message) == TableResolverStatus:
+          finished_table_resovlers += 1
+          if total_table_resolvers == finished_table_resovlers:       
+            break
+          continue
         yield message
       thread.shutdown()
diff --git a/cloudquery/sdk/scheduler/table_resolver.py b/cloudquery/sdk/scheduler/table_resolver.py
@@ -1,20 +1,20 @@
 
 from cloudquery.sdk.schema.table import Table
 from cloudquery.sdk.schema import Resource
-from typing import Any
+from typing import Any,Generator
 
 class TableResolver:
     def __init__(self, table: Table) -> None:
         self._table = table
     
     @property
-    def table(self):
+    def table(self) -> Table:
       return self._table
 
     def multiplex(self, client):
       return [client]
     
-    def resolve(self, client, parent_resource) -> Any:
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
       raise NotImplementedError()
 
     def pre_resource_resolve(self, client, resource):
diff --git a/cloudquery/sdk/schema/resource.py b/cloudquery/sdk/schema/resource.py
@@ -1,16 +1,28 @@
 
 from .table import Table
 from typing import Any
+import pyarrow as pa
 from cloudquery.sdk.scalar import ScalarFactory
 
 class Resource:
     def __init__(self, table: Table, parent, item: Any) -> None:
       self._table = table
       self._parent = parent
+      self._item = item
       factory = ScalarFactory()
       self._data = [factory.new_scalar(i.type) for i in table.columns]
     
+    @property
+    def item(self):
+      return self._item
+
     def set(self, column_name: str, value: Any):
       index_column = self._table.index_column(column_name)
       self._data[index_column].set(value)
+    
+    def to_list_of_arr(self):
+       return [[self._data[i].value] for i, _ in enumerate(self._table.columns)]
+    
+    def to_arrow_record(self):
+       return pa.record_batch(self.to_list_of_arr(), schema=self._table.to_arrow_schema())
     
diff --git a/tests/scheduler/__init__.py b/tests/scheduler/__init__.py
diff --git a/tests/scheduler/scheduler.py b/tests/scheduler/scheduler.py
@@ -0,0 +1,33 @@
+
+from typing import Any, List, Generator
+import pyarrow as pa
+from cloudquery.sdk.scheduler import Scheduler, TableResolver
+from cloudquery.sdk.schema import Table, Column, Resource
+from cloudquery.sdk.message import SyncMessage
+from cloudquery.sdk.schema.table import Table
+
+class SchedulerTestTable(Table):
+    def __init__(self):
+        super().__init__("test_table", [
+            Column("test_column", pa.int64())
+        ])
+
+class SchedulerTestTableResolver(TableResolver):
+    def __init__(self) -> None:
+        super().__init__(SchedulerTestTable())
+    
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+      yield {"test_column": 1}
+
+class TestClient:
+    pass
+
+def test_scheduler():
+    client = TestClient()
+    s = Scheduler(10)
+    resources: List[SyncMessage] = []
+    for resource in s.sync(client, [SchedulerTestTableResolver()]):
+        resources.append(resource)
+    assert len(resources) == 2
+    print(resources[1].record.to_pydict())
+    
diff --git a/tests/schema/test_column.py b/tests/schema/test_column.py
@@ -3,7 +3,7 @@
 
 
 def test_column():
-    c = Column("test_column", pa.int32())
+    c = Column("test_column", pa.int64())
     f = c.to_arrow_field()
     c1 = Column.from_arrow_field(f)
     f1 = c1.to_arrow_field()

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+from .scheduler import Scheduler`
	`2`	`+from .table_resolver import TableResolver`