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

cloudquery/sdk/plugin/plugin.py

@@ -0,0 +1,102 @@
+from dataclasses import dataclass
+from typing import List, Generator
+
+from cloudquery.sdk import message
+from cloudquery.sdk.schema import Table
+
+MIGRATE_MODE_STRINGS = ["safe", "force"]
+
+
+@dataclass
+class TableOptions:
+    tables: List[str] = None
+    skip_tables: List[str] = None
+    skip_dependent_tables: bool = False
+
+
+@dataclass
+class BackendOptions:
+    connection: str = None
+    table_name: str = None
+    # set when the CLI persists state sent as checkpoints on the sync stream
+    supports_checkpoints: bool = False
+
+
+@dataclass
+class SyncOptions:
+    deterministic_cq_id: bool = False
+    skip_dependent_tables: bool = False
+    skip_tables: List[str] = None
+    tables: List[str] = None
+    backend_options: BackendOptions = None
+
+
+@dataclass
+class BuildTarget:
+    os: str = None
+    arch: str = None
+
+
+@dataclass
+class Options:
+    dockerfile: str = None
+    build_targets: List[BuildTarget] = None
+    team: str = None
+    kind: str = None
+    json_schema: str = None
+
+
+class Plugin:
+    def __init__(self, name: str, version: str, opts: Options = None) -> None:
+        self._name = name
+        self._version = version
+        self._opts = Options() if opts is None else opts
+        if self._opts.dockerfile is None:
+            self._opts.dockerfile = "Dockerfile"
+        if self._opts.build_targets is None:
+            self._opts.build_targets = [
+                BuildTarget("linux", "amd64"),
+                BuildTarget("linux", "arm64"),
+            ]
+
+    def init(self, spec: bytes, no_connection: bool = False) -> None:
+        pass
+
+    def set_logger(self, logger) -> None:
+        pass
+
+    def name(self) -> str:
+        return self._name
+
+    def version(self) -> str:
+        return self._version
+
+    def team(self) -> str:
+        return self._opts.team
+
+    def kind(self) -> str:
+        return self._opts.kind
+
+    def json_schema(self) -> str:
+        return self._opts.json_schema
+
+    def dockerfile(self) -> str:
+        return self._opts.dockerfile
+
+    def build_targets(self) -> List[BuildTarget]:
+        return self._opts.build_targets
+
+    def get_tables(self, options: TableOptions) -> List[Table]:
+        raise NotImplementedError()
+
+    def sync(self, options: SyncOptions) -> Generator[message.SyncMessage, None, None]:
+        raise NotImplementedError()
+
+    def write(self, writer: Generator[message.WriteMessage, None, None]) -> None:
+        raise NotImplementedError()
+
+    def read(self, table: Table) -> Generator[message.ReadMessage, None, None]:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        raise NotImplementedError()

cloudquery/sdk/scalar/__init__.py

@@ -0,0 +1,15 @@
+from .scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
+from .scalar_factory import ScalarFactory
+from .vector import Vector
+
+from .binary import Binary
+from .bool import Bool
+from .date32 import Date32
+from .date64 import Date64
+from .float import Float
+from .int import Int
+from .list import List
+from .string import String
+from .timestamp import Timestamp
+from .uint import Uint
+from .uuid import UUID

cloudquery/sdk/scalar/binary.py

@@ -0,0 +1,35 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from .scalar import NULL_VALUE
+
+
+class Binary(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, Binary):
+            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
+
+        if isinstance(value, Binary):
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, bytes):
+            self._valid = True
+            self._value = value
+        elif isinstance(value, str):
+            self._valid = True
+            self._value = value.encode()
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for Binary scalar"
+            )

cloudquery/sdk/scalar/bool.py

@@ -0,0 +1,48 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
+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
+    if lower_input in false_strings:
+        return False
+    else:
+        raise ScalarInvalidTypeError(f"Invalid boolean string: {input_string}")
+
+
+class Bool(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(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:
+            self._valid = False
+            return
+
+        if isinstance(value, Bool):
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, bool):
+            self._value = value
+        elif isinstance(value, str):
+            self._value = parse_string_to_bool(value)
+        else:
+            raise ScalarInvalidTypeError(f"Invalid type {type(value)} for Bool scalar")
+
+        self._valid = True

cloudquery/sdk/scalar/date32.py

@@ -0,0 +1,39 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
+from datetime import datetime, time
+from typing import Any
+
+
+class Date32(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(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:
+            self._valid = False
+            return
+
+        if isinstance(value, Date32):
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, datetime):
+            self._value = value
+        elif isinstance(value, str):
+            self._value = datetime.strptime(value, "%Y-%m-%d")
+        elif isinstance(value, time):
+            self._value = datetime.combine(datetime.today(), value)
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for Date32 scalar"
+            )
+
+        self._valid = True

cloudquery/sdk/scalar/date64.py

@@ -0,0 +1,39 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
+from datetime import datetime, time
+from typing import Any
+
+
+class Date64(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, Date64):
+            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, Date64):
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, datetime):
+            self._value = value
+        elif isinstance(value, str):
+            self._value = datetime.strptime(value, "%Y-%m-%d")
+        elif isinstance(value, time):
+            self._value = datetime.combine(datetime.today(), value)
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for Date64 scalar"
+            )
+
+        self._valid = True

cloudquery/sdk/scalar/float.py

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

cloudquery/sdk/scalar/int.py

@@ -0,0 +1,56 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+
+class Int(Scalar):
+    def __init__(self, valid: bool = False, value: any = None, bitwidth: int = 64):
+        self._bitwidth = bitwidth
+        self._min = -(2 ** (bitwidth - 1))
+        self._max = 2 ** (bitwidth - 1)
+        super().__init__(valid, value)
+
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, Int) and self._bitwidth == scalar.bitwidth:
+            return 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, Int) 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 type for Int{self._bitwidth} scalar"
+                ) from e
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for Int{self._bitwidth} scalar"
+            )
+        if val < self._min or val >= self._max:
+            raise ScalarInvalidTypeError(
+                f"Invalid Int{self._bitwidth} scalar with value {val}"
+            )
+        self._value = val
+        self._valid = True

cloudquery/sdk/scalar/json.py

@@ -0,0 +1,27 @@
+import json
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+
+class JSON(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, JSON):
+            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
+
+        if isinstance(value, (str, bytes)):
+            # test if it is a valid json
+            json.loads(value)
+            self._value = value
+        else:
+            self._value = json.dumps(value)
+        self._valid = True

cloudquery/sdk/scalar/list.py

@@ -0,0 +1,56 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from .scalar import NULL_VALUE
+from .vector import Vector
+from typing import Any, Type
+
+
+class List(Scalar):
+    def __init__(self, scalar_type: Type[Scalar]):
+        super().__init__(False, None)
+        self._value = Vector(scalar_type)
+        self._type = scalar_type
+
+    def __eq__(self, other: "List") -> bool:
+        if (not isinstance(other, self.__class__)) or self._valid != other._valid:
+            return False
+        return self._value == other._value
+
+    @property
+    def type(self):
+        return self._type
+
+    @property
+    def value(self):
+        return self._value
+
+    def set(self, value: Any):
+        if value is None:
+            self._valid = False
+            self._value = Vector()
+            return
+
+        if isinstance(value, self._type):
+            if not value.is_valid:
+                self._valid = False
+                self._value = Vector()
+                return
+            return self.set([value.value])
+
+        if isinstance(value, (list, tuple)):
+            self._value = Vector()
+            for item in value:
+                scalar = self._type()
+                scalar.set(item)
+                self._value.append(scalar)
+                self._valid = True
+            return
+
+        raise ScalarInvalidTypeError(f"Invalid type {type(value)} for List")
+
+    def __str__(self) -> str:
+        if not self._valid:
+            return NULL_VALUE
+        return f"[{', '.join(str(v) for v in self._value)}]"
+
+    def __len__(self):
+        return len(self.value.data)

cloudquery/sdk/scalar/scalar.py

@@ -0,0 +1,30 @@
+from abc import abstractmethod
+
+NULL_VALUE = "null"
+
+
+class ScalarInvalidTypeError(Exception):
+    pass
+
+
+class Scalar:
+    def __init__(self, valid: bool = False, value: any = None):
+        self._valid = valid
+        self._value = None
+        self.set(value)
+
+    @property
+    def is_valid(self) -> bool:
+        return self._valid
+
+    def __str__(self) -> str:
+        return str(self._value) if self._valid else NULL_VALUE
+
+    @property
+    @abstractmethod
+    def value(self):
+        pass
+
+    @abstractmethod
+    def set(self, value: any):
+        pass

cloudquery/sdk/scalar/scalar_factory.py

@@ -0,0 +1,64 @@
+from functools import partial
+
+import pyarrow as pa
+
+from cloudquery.sdk.types import UUIDType, JSONType
+from .binary import Binary
+from .bool import Bool
+from .date32 import Date32
+from .date64 import Date64
+from .float import Float
+from .int import Int
+from .json import JSON
+from .list import List
+from .scalar import ScalarInvalidTypeError
+from .string import String
+from .timestamp import Timestamp
+from .uint import Uint
+from .uuid import UUID
+
+
+class ScalarFactory:
+    def __init__(self):
+        pass
+
+    def new_scalar(self, dt: pa.DataType):
+        dt_id = dt.id
+        type_id__map = {
+            pa.types.lib.Type_INT64: partial(Int, bitwidth=64),
+            pa.types.lib.Type_INT32: partial(Int, bitwidth=32),
+            pa.types.lib.Type_INT16: partial(Int, bitwidth=16),
+            pa.types.lib.Type_INT8: partial(Int, bitwidth=8),
+            pa.types.lib.Type_UINT64: partial(Uint, bitwidth=64),
+            pa.types.lib.Type_UINT32: partial(Uint, bitwidth=32),
+            pa.types.lib.Type_UINT16: partial(Uint, bitwidth=16),
+            pa.types.lib.Type_UINT8: partial(Uint, bitwidth=8),
+            pa.types.lib.Type_BINARY: Binary,
+            pa.types.lib.Type_LARGE_BINARY: Binary,
+            pa.types.lib.Type_FIXED_SIZE_BINARY: Binary,
+            pa.types.lib.Type_BOOL: Bool,
+            pa.types.lib.Type_DATE64: Date64,
+            pa.types.lib.Type_DATE32: Date32,
+            pa.types.lib.Type_DOUBLE: partial(Float, bitwidth=64),
+            pa.types.lib.Type_FLOAT: partial(Float, bitwidth=32),
+            pa.types.lib.Type_HALF_FLOAT: partial(Float, bitwidth=16),
+            pa.types.lib.Type_LIST: List,
+            pa.types.lib.Type_LARGE_LIST: List,
+            pa.types.lib.Type_FIXED_SIZE_LIST: List,
+            pa.types.lib.Type_STRING: String,
+            pa.types.lib.Type_LARGE_STRING: String,
+            pa.types.lib.Type_TIMESTAMP: Timestamp,
+        }
+        # Built-in Types
+        if dt_id in type_id__map:
+            scalar_type = type_id__map[dt_id]
+            if scalar_type == List:
+                item = self.new_scalar(dt.field(0).type)
+                return scalar_type(type(item))
+            return scalar_type()
+        # Extension Types - Can't do the same trick as above as they don't have `id`s and they are not hashable. :(
+        if dt == UUIDType():
+            return UUID()
+        if dt == JSONType():
+            return JSON()
+        raise ScalarInvalidTypeError(f"Invalid type {dt} for scalar")

cloudquery/sdk/scalar/string.py

@@ -0,0 +1,31 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+
+
+class String(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, String):
+            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
+
+        if isinstance(value, String):
+            self._valid = value._valid
+            self._value = value.value
+            return
+
+        if isinstance(value, str):
+            self._valid = True
+            self._value = value
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for String scalar"
+            )

cloudquery/sdk/scalar/timestamp.py

@@ -0,0 +1,41 @@
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
+from datetime import datetime, time
+from typing import Any
+import pandas as pd
+
+
+class Timestamp(Scalar):
+    def __eq__(self, scalar: Scalar) -> bool:
+        if scalar is None:
+            return False
+        if isinstance(scalar, Timestamp):
+            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
+
+        if isinstance(value, Timestamp):
+            self._valid = value.is_valid
+            self._value = value.value
+            return
+
+        if isinstance(value, pd.Timestamp):
+            self._value = value
+        elif isinstance(value, datetime):
+            self._value = pd.to_datetime(value)
+        elif isinstance(value, str):
+            self._value = pd.to_datetime(value)
+        elif isinstance(value, time):
+            self._value = pd.to_datetime(datetime.combine(datetime.today(), value))
+        else:
+            raise ScalarInvalidTypeError(
+                f"Invalid type {type(value)} for Timestamp scalar"
+            )
+
+        self._valid = True