more work on scalar package

Author: yevgenypats

cloudquery/sdk/scalar/__init__.py

@@ -1,7 +1,8 @@
-from .scalar import Scalar, ScalarInvalidTypeError
+from .scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
 from .scalar_factory import ScalarFactory
 from .binary import Binary
 from .bool import Bool
 from .date32 import Date32
 from .float64 import Float64
-from .int64 import Int64
+from .int64 import Int64
+from .uuid import UUID

cloudquery/sdk/scalar/binary.py

@@ -1,4 +1,5 @@
 from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from .scalar import NULL_VALUE
 
 class Binary(Scalar):
     def __init__(self, valid: bool = False, value: bytes = None):
@@ -11,6 +12,9 @@ def __eq__(self, scalar: Scalar) -> bool:
       if type(scalar) == Binary:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
+    
+    def __str__(self) -> str:
+       return str(self._value) if self._valid else NULL_VALUE
 
     @property
     def value(self):

cloudquery/sdk/scalar/bool.py

@@ -1,5 +1,5 @@
 
-from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
 from typing import Any
 
 def parse_string_to_bool(input_string):
@@ -27,6 +27,9 @@ def __eq__(self, scalar: Scalar) -> bool:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
 
+    def __str__(self) -> str:
+       return str(self._value) if self._valid else NULL_VALUE
+
     @property
     def value(self):
        return self._value

cloudquery/sdk/scalar/date32.py

@@ -1,5 +1,5 @@
 
-from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError
+from cloudquery.sdk.scalar import Scalar, ScalarInvalidTypeError, NULL_VALUE
 from datetime import datetime, time
 from typing import Any
 
@@ -15,6 +15,9 @@ def __eq__(self, scalar: Scalar) -> bool:
           return self._value == scalar._value and self._valid == scalar._valid
       return False
 
+    def __str__(self) -> str:
+       return str(self._value) if self._valid else NULL_VALUE
+    
     @property
     def value(self):
        return self._value

cloudquery/sdk/scalar/scalar.py

@@ -1,4 +1,6 @@
 
+NULL_VALUE = "null"
+
 class ScalarInvalidTypeError(Exception):
   pass
 

cloudquery/sdk/scalar/uuid.py

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

cloudquery/sdk/types/uuid.py

@@ -1,9 +1,21 @@
 import pyarrow as pa
+import pyarrow
+import sys
 
 class UuidType(pa.PyExtensionType):
-
     def __init__(self):
         pa.PyExtensionType.__init__(self, pa.binary(16))
 
     def __reduce__(self):
-        return UuidType, ()
+        return UuidType, ()
+
+    def __arrow_ext_serialize__(self):
+      # since we don't have a parameterized type, we don't need extra
+      # metadata to be deserialized
+      return b'uuid-serialized'
+
+    @classmethod
+    def __arrow_ext_deserialize__(self, storage_type, serialized):
+      # return an instance of this subclass given the serialized
+      # metadata.
+      return UuidType()

tests/scalar/uuid.py

@@ -0,0 +1,11 @@
+import pytest
+import uuid
+from cloudquery.sdk.scalar import UUID
+
+@pytest.mark.parametrize("input_value,expected_scalar", [
+    ("550e8400-e29b-41d4-a716-446655440000", UUID(True, uuid.UUID("550e8400-e29b-41d4-a716-446655440000"))),
+])
+def test_binary_set(input_value, expected_scalar):
+    b = UUID()
+    b.set(input_value)
+    assert b == expected_scalar