cloudquery-plugin-sdk 0.1.55__py3-none-any.whl

cloudquery/sdk/scalar/uint.py

@@ -0,0 +1,58 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+
+class Uint(Scalar):
+    def __init__(self, valid: bool = False, value: any = None, bitwidth: int = 64):
+        self._bitwidth = bitwidth
+        self._max = 2**bitwidth
+        super().__init__(valid, value)
+
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, Uint):
+            return (
+                self._bitwidth == scalar.bitwidth
+                and self._value == scalar._value
+                and self._valid == scalar._valid
+            )
+        return False
+
+    @property
+    def bitwidth(self):
+        return self._bitwidth
+
+    @property
+    def value(self):
+        return self._value
+
+    def set(self, value: any):
+        if value is None:
+            self._valid = False
+            return
+
+        if isinstance(value, Uint) and value.bitwidth == self._bitwidth:
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, int):
+            val = value
+        elif isinstance(value, float):
+            val = int(value)
+        elif isinstance(value, str):
+            try:
+                val = int(value)
+            except ValueError as e:
+                raise ScalarInvalidTypeError(
+                    f"Invalid value for Int{self._bitwidth} scalar"
+                ) from e
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for Int{self._bitwidth} scalar"
+            )
+
+        if val < 0 or val >= self._max:
+            raise ScalarInvalidTypeError(f"Invalid Uint{self._bitwidth} scalar {val}")
+        self._value = val
+        self._valid = True

cloudquery/sdk/scalar/uuid.py

@@ -0,0 +1,43 @@
+import uuid
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+
+class UUID(Scalar):
+    def __init__(self, valid: bool = False, value: uuid.UUID = None):
+        if isinstance(value, uuid.UUID):
+            value = value.bytes
+        super().__init__(valid, value)
+
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, UUID):
+            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:
+            self._valid = False
+            return
+
+        if isinstance(value, UUID):
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, bytes):
+            self._value = value
+        elif isinstance(value, uuid.UUID):
+            self._value = value.bytes
+        elif isinstance(value, str):
+            try:
+                self._value = uuid.UUID(value).bytes
+            except ValueError as e:
+                raise ScalarInvalidTypeError("Invalid type for UUID scalar") from e
+        else:
+            raise ScalarInvalidTypeError(f"Invalid type {type(value)} for UUID scalar")
+        self._valid = True

cloudquery/sdk/scalar/vector.py

@@ -0,0 +1,43 @@
+from .scalar import Scalar
+from typing import Type
+
+
+class Vector:
+    def __init__(self, type: Type[Scalar] = None, *args):
+        self.data = []
+        self.type = type
+        for arg in args:
+            self.append(arg)  # Use the append method for type checking and appending
+
+    def append(self, item):
+        if not isinstance(item, Scalar):
+            raise TypeError("Item is not of type Scalar or its subclass")
+
+        if self.type is None:
+            self.type = type(item)
+            self.data.append(item)
+        elif isinstance(item, self.type):
+            self.data.append(item)
+        else:
+            raise TypeError(f"Item is not of type {self.type.__name__}")
+
+    def __eq__(self, other):
+        if not isinstance(other, Vector):
+            return False
+        if len(self) != len(other):
+            return False
+
+        for self_item, other_item in zip(self.data, other.data):
+            if self_item != other_item:
+                return False
+
+        return True
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def __len__(self):
+        return len(self.data)
+
+    def __repr__(self):
+        return f"Vector of {self.type.__name__ if self.type else 'unknown type'}: {self.data}"

cloudquery/sdk/scheduler/__init__.py

@@ -0,0 +1,3 @@
+from .scheduler import Scheduler
+from .table_resolver import TableResolver
+from .table_resolver import Client

cloudquery/sdk/scheduler/scheduler.py

@@ -0,0 +1,267 @@
+import queue
+from concurrent import futures
+from typing import List, Generator, Any, Optional
+
+import structlog
+
+from cloudquery.sdk.message import (
+    SyncMessage,
+    SyncInsertMessage,
+    SyncMigrateTableMessage,
+    SyncErrorMessage,
+)
+from cloudquery.sdk.internal.checkpoint import Checkpoints, Scope
+from cloudquery.sdk.schema import Resource
+from cloudquery.sdk.stateclient.stateclient import StateClient
+from .table_resolver import TableResolver, Client
+
+QUEUE_PER_WORKER = 100
+
+
+class ThreadPoolExecutorWithQueueSizeLimit(futures.ThreadPoolExecutor):
+    def __init__(self, maxsize, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._work_queue = queue.Queue(maxsize=maxsize)
+
+
+class TableResolverStarted:
+    def __init__(self) -> None:
+        pass
+
+
+class TableResolverFinished:
+    def __init__(self) -> None:
+        pass
+
+
+class TableResolversSubmitted:
+    def __init__(self) -> None:
+        pass
+
+
+class Scheduler:
+    def __init__(
+        self, concurrency: int, queue_size: int = 0, max_depth: int = 3, logger=None
+    ):
+        self._post_sync_hook = lambda: None
+        self._queue = queue.Queue()
+        # pass to StateClientBuilder.build so that state flushes are sent as checkpoints, when the CLI supports them
+        self.checkpoints = Checkpoints()
+        self._max_depth = max_depth
+        if logger is None:
+            self._logger = structlog.get_logger()
+        else:
+            self._logger = logger
+        if concurrency <= 0:
+            raise ValueError("concurrency 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: List[ThreadPoolExecutorWithQueueSizeLimit] = []
+        current_depth_concurrency = concurrency
+        current_depth_queue_size = queue_size
+        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 shutdown(self):
+        for pool in self._pools:
+            pool.shutdown()
+
+    def resolve_resource(
+        self, resolver: TableResolver, client: Client, parent: Resource, item: Any
+    ) -> Resource:
+        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,
+        depth: int,
+        client: Client,
+        parent_item: Optional[Resource],
+        res: queue.Queue,
+        scope: Optional[Scope] = None,
+    ):
+        scope = scope or self.checkpoints.new_scope(None)
+        try:
+            with self.checkpoints.entered(scope):
+                self._resolve_table(resolver, depth, client, parent_item, res, scope)
+        finally:
+            # after every resource of the resolver is on the queue, so its checkpoints follow them
+            self.checkpoints.resolver_finished(scope)
+            res.put(TableResolverFinished())
+
+    def _resolve_table(
+        self,
+        resolver: TableResolver,
+        depth: int,
+        client: Client,
+        parent_item: Optional[Resource],
+        res: queue.Queue,
+        scope: Scope,
+    ):
+        total_resources = 0
+        try:
+            if depth == 0:
+                self._logger.info(
+                    "table resolver started",
+                    client_id=client.id(),
+                    table=resolver.table.name,
+                    depth=depth,
+                )
+            else:
+                self._logger.debug(
+                    "table resolver started",
+                    client_id=client.id(),
+                    table=resolver.table.name,
+                    depth=depth,
+                )
+            for item in resolver.resolve(client, parent_item):
+                try:
+                    resource = self.resolve_resource(
+                        resolver, client, parent_item, item
+                    )
+                except Exception as e:
+                    self._logger.error(
+                        "failed to resolve resource",
+                        client_id=client.id(),
+                        table=resolver.table.name,
+                        depth=depth,
+                        exc_info=e,
+                    )
+                    self._logger.debug(
+                        "details about resource that failed to resolve",
+                        client_id=client.id(),
+                        table=resolver.table.name,
+                        resource=item,
+                    )
+                    continue
+                res.put(SyncInsertMessage(resource.to_arrow_record()))
+                for child_resolvers in resolver.child_resolvers:
+                    res.put(TableResolverStarted())
+                    self._pools[depth + 1].submit(
+                        self.resolve_table,
+                        child_resolvers,
+                        depth + 1,
+                        client,
+                        resource,
+                        res,
+                        self.checkpoints.new_scope(scope),
+                    )
+                total_resources += 1
+            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=e,
+            )
+            res.put(SyncErrorMessage(resolver.table.name, str(e)))
+
+    def _sync(
+        self,
+        client,
+        resolvers: List[TableResolver],
+        res: queue.Queue,
+        deterministic_cq_id=False,
+    ):
+        try:
+            for resolver in resolvers:
+                clients = resolver.multiplex(client)
+                for client in clients:
+                    res.put(TableResolverStarted())
+                    self._pools[0].submit(
+                        self.resolve_table, resolver, 0, client, None, res
+                    )
+        finally:
+            res.put(TableResolversSubmitted())
+
+    def sync(
+        self, client, resolvers: List[TableResolver], deterministic_cq_id=False
+    ) -> Generator[SyncMessage, None, None]:
+        res = queue.Queue()
+        yield from self._send_migrate_table_messages(resolvers)
+
+        self.checkpoints.start(res)
+        try:
+            thread = futures.ThreadPoolExecutor()
+            thread.submit(self._sync, client, resolvers, res, deterministic_cq_id)
+            submitted = False
+            total_table_resolvers = 0
+            finished_table_resolvers = 0
+            while not (submitted and total_table_resolvers == finished_table_resolvers):
+                message = res.get()
+                if isinstance(message, TableResolversSubmitted):
+                    submitted = True
+                elif isinstance(message, TableResolverStarted):
+                    total_table_resolvers += 1
+                elif isinstance(message, TableResolverFinished):
+                    finished_table_resolvers += 1
+                else:
+                    yield message
+
+            self._post_sync_hook()
+            # checkpoints the post-sync hook flushed, after every resource
+            while not res.empty():
+                yield res.get_nowait()
+            thread.shutdown(wait=True)
+        finally:
+            self.checkpoints.stop()
+
+    def _send_migrate_table_messages(
+        self, resolvers: List[TableResolver]
+    ) -> Generator[SyncMessage, None, None]:
+        for resolver in resolvers:
+            yield SyncMigrateTableMessage(table=resolver.table.to_arrow_schema())
+            if resolver.child_resolvers:
+                yield from self._send_migrate_table_messages(resolver.child_resolvers)
+
+    def set_post_sync_hook(self, fn):
+        """
+        Use this to set a function that will be called after the sync is finished,
+        a la `defer fn()` in Go (but for a single function, rather than a stack).
+
+        This is necessary because plugins use this pattern on their sync method:
+
+        ```
+        return self._scheduler.sync(...)
+        ```
+
+        So if a plugin has a `state_client`, there's nowhere to call the flush method.
+        """
+        self._post_sync_hook = fn

cloudquery/sdk/scheduler/table_resolver.py

@@ -0,0 +1,48 @@
+from cloudquery.sdk.schema.table import Table
+from cloudquery.sdk.schema import Resource
+from typing import Any, Generator, List, Optional
+
+
+class Client:
+    def id(self) -> str:
+        raise NotImplementedError()
+
+
+class TableResolver:
+    def __init__(self, table: Table, child_resolvers: Optional[List] = None) -> None:
+        if child_resolvers is None:
+            child_resolvers = []
+        self._table = table
+        self._child_resolvers = child_resolvers
+        self.state_client = None
+
+    @property
+    def table(self) -> Table:
+        return self._table
+
+    @property
+    def child_resolvers(self):
+        return self._child_resolvers
+
+    def set_state_client(self, state_client):
+        self.state_client = state_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: Client, resource):
+        return
+
+    def resolve_column(self, client: Client, resource: Resource, column_name: str):
+        if isinstance(resource.item, 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, getattr(resource.item, column_name))
+
+    def post_resource_resolve(self, client: Client, resource):
+        return

cloudquery/sdk/schema/__init__.py

@@ -0,0 +1,16 @@
+from .column import Column
+from .table import (
+    Table,
+    tables_to_arrow_schemas,
+    filter_dfs,
+    TableColumnChangeType,
+    TableColumnChange,
+    TableColumnChangeType,
+    get_table_changes,
+    get_table_column,
+    flatten_tables_recursive,
+    flatten_tables,
+)
+from .resource import Resource
+
+# from .table_resolver import TableReso

cloudquery/sdk/schema/arrow.py

@@ -0,0 +1,11 @@
+METADATA_UNIQUE = b"cq:extension:unique"
+METADATA_PRIMARY_KEY = b"cq:extension:primary_key"
+METADATA_CONSTRAINT_NAME = b"cq:extension:constraint_name"
+METADATA_INCREMENTAL = b"cq:extension:incremental"
+
+METADATA_TRUE = b"true"
+METADATA_FALSE = b"false"
+METADATA_TABLE_NAME = b"cq:table_name"
+METADATA_TABLE_DESCRIPTION = b"cq:table_description"
+METADATA_TABLE_TITLE = b"cq:table_title"
+METADATA_TABLE_DEPENDS_ON = b"cq:table_depends_on"

cloudquery/sdk/schema/column.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import pyarrow as pa
+
+from cloudquery.sdk.schema import arrow
+
+
+class Column:
+    def __init__(
+        self,
+        name: str,
+        type: pa.DataType,
+        description: str = "",
+        primary_key: bool = False,
+        not_null: bool = False,
+        incremental_key: bool = False,
+        unique: bool = False,
+    ) -> None:
+        self.name = name
+        self.type = type
+        self.description = description
+        self.primary_key = primary_key
+        self.not_null = not_null
+        self.incremental_key = incremental_key
+        self.unique = unique
+
+    def __str__(self) -> str:
+        return f"Column(name={self.name}, type={self.type}, description={self.description}, primary_key={self.primary_key}, not_null={self.not_null}, incremental_key={self.incremental_key}, unique={self.unique})"
+
+    def __repr__(self) -> str:
+        return f"Column(name={self.name}, type={self.type}, description={self.description}, primary_key={self.primary_key}, not_null={self.not_null}, incremental_key={self.incremental_key}, unique={self.unique})"
+
+    def __eq__(self, __value: object) -> bool:
+        if isinstance(__value, Column):
+            return (
+                self.name == __value.name
+                and self.type == __value.type
+                and self.description == __value.description
+                and self.primary_key == __value.primary_key
+                and self.not_null == __value.not_null
+                and self.incremental_key == __value.incremental_key
+                and self.unique == __value.unique
+            )
+        return False
+
+    def to_arrow_field(self):
+        metadata = {
+            arrow.METADATA_PRIMARY_KEY: (
+                arrow.METADATA_TRUE if self.primary_key else arrow.METADATA_FALSE
+            ),
+            arrow.METADATA_UNIQUE: (
+                arrow.METADATA_TRUE if self.unique else arrow.METADATA_FALSE
+            ),
+            arrow.METADATA_INCREMENTAL: (
+                arrow.METADATA_TRUE if self.incremental_key else arrow.METADATA_FALSE
+            ),
+        }
+        return pa.field(
+            self.name, self.type, metadata=metadata, nullable=not self.not_null
+        )
+
+    @staticmethod
+    def from_arrow_field(field: pa.Field) -> Column:
+        metadata = field.metadata or {}
+        primary_key = metadata.get(arrow.METADATA_PRIMARY_KEY) == arrow.METADATA_TRUE
+        unique = metadata.get(arrow.METADATA_UNIQUE) == arrow.METADATA_TRUE
+        incremental_key = (
+            metadata.get(arrow.METADATA_INCREMENTAL) == arrow.METADATA_TRUE
+        )
+        return Column(
+            field.name,
+            field.type,
+            primary_key=primary_key,
+            not_null=not field.nullable,
+            unique=unique,
+            incremental_key=incremental_key,
+        )

cloudquery/sdk/schema/resource.py

@@ -0,0 +1,29 @@
+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()
+        )