spinta 0.2.dev23__py3-none-any.whl → 0.2.dev25__py3-none-any.whl

spinta/backends/__init__.py

@@ -13,6 +13,7 @@ from typing import Iterable
 from typing import List
 from typing import Optional
 
+from cbor2 import dumps as cbor_dumps
 import dateutil
 import shapely.geometry.base
 from geoalchemy2.elements import WKTElement, WKBElement
@@ -23,7 +24,7 @@ from spinta import commands
 from spinta import exceptions
 from spinta.backends.components import Backend
 from spinta.backends.components import SelectTree
-from spinta.backends.helpers import check_unknown_props, get_select_tree, prepare_response
+from spinta.backends.helpers import check_unknown_props, get_select_tree, prepare_response, is_accessible_by_equals_sign
 from spinta.backends.helpers import flat_select_to_nested
 from spinta.backends.helpers import get_model_reserved_props
 from spinta.backends.helpers import get_select_prop_names
@@ -58,7 +59,19 @@ from spinta.exceptions import (
 from spinta.exceptions import NoItemRevision
 from spinta.formats.components import Format
 from spinta.manifests.components import Manifest
-from spinta.types.datatype import Array, ExternalRef, Inherit, PageType, BackRef, ArrayBackRef, Integer, Boolean, Denorm
+from spinta.types.datatype import (
+    Array,
+    ExternalRef,
+    Inherit,
+    PageType,
+    BackRef,
+    ArrayBackRef,
+    Integer,
+    Boolean,
+    Denorm,
+    Base32,
+    String,
+)
 from spinta.types.datatype import Binary
 from spinta.types.datatype import DataType
 from spinta.types.datatype import Date
@@ -590,12 +603,31 @@ def is_object_id(context: Context, value: str):
 
 @is_object_id.register(Context, Backend, Model, str)
 def is_object_id(context: Context, backend: Backend, model: Model, value: str):
+    return is_object_id(context, backend, model.properties["_id"].dtype, value)
+
+
+@is_object_id.register(Context, Backend, PrimaryKey, str)
+def is_object_id(context: Context, backend: Backend, dtype: PrimaryKey, value: str):
     try:
         return uuid.UUID(value).version == 4
     except ValueError:
         return False
 
 
+@is_object_id.register(Context, Backend, DataType, str)
+def is_object_id(context: Context, backend: Backend, dtype: DataType, value: str):
+    candidate = value
+    if is_accessible_by_equals_sign(dtype.prop, value):
+        if not value.startswith("="):
+            return False
+        candidate = value[1:]
+    try:
+        dtype.load(candidate)
+    except exceptions.InvalidValue:
+        return False
+    return True
+
+
 @is_object_id.register(Context, Backend, Model, uuid.UUID)
 def is_object_id(context: Context, backend: Backend, model: Model, value: uuid.UUID):
     return value.version == 4
@@ -1750,6 +1782,13 @@ def cast_backend_to_python(context: Context, dtype: DataType, backend: Backend,
     return data
 
 
+@commands.cast_backend_to_python.register(Context, String, Backend, object)
+def cast_backend_to_python(context: Context, dtype: String, backend: Backend, data: Any, **kwargs) -> Any:
+    if data is None or is_nan(data):
+        return None
+    return str(data)
+
+
 @commands.cast_backend_to_python.register(Context, UUID, Backend, object)
 def cast_backend_to_python(context: Context, dtype: UUID, backend: Backend, data: Any, **kwargs) -> Any:
     if is_nan(data):
@@ -1891,6 +1930,13 @@ def cast_backend_to_python(context: Context, dtype: Ref, backend: Backend, data:
 
     processed_data = {}
     for key in data:
+        if key == "_id":
+            # _id reaches this dispatch already in its final form — produced by
+            # handle_ref_key_assignment for external readers, or read directly
+            # from the storage column for internal backends. Re-applying the
+            # referenced model's _id cast double-encodes Base32 ids.
+            processed_data[key] = data[key]
+            continue
         prop = commands.resolve_property(dtype.prop.model, f"{dtype.prop.place}.{key}")
         if prop is not None:
             processed_data[key] = commands.cast_backend_to_python(context, prop, backend, data[key], **kwargs)
@@ -1939,6 +1985,18 @@ def cast_backend_to_python(context: Context, dtype: Denorm, backend: Backend, da
     return commands.cast_backend_to_python(context, dtype.rel_prop, backend, data, **kwargs)
 
 
+@commands.cast_backend_to_python.register(Context, Base32, Backend, object)
+def cast_backend_to_python(context: Context, dtype: Base32, backend: Backend, data: Any, **kwargs) -> Any:
+    if is_nan(data):
+        return None
+    if isinstance(data, (list, tuple)):
+        data = cbor_dumps(list(data))
+    else:
+        data = str(data).encode("utf-8")
+    encoded = base64.b32encode(data)
+    return encoded.rstrip(b"=").decode("utf-8")
+
+
 @commands.reload_backend_metadata.register(Context, Manifest, Backend)
 def reload_backend_metadata(context, manifest, backend):
     pass

spinta/backends/components.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import contextlib
+import dataclasses
 from typing import Any, Type
 from typing import Dict
 from typing import Optional
 from typing import Set
 
-from spinta.backends.constants import BackendOrigin, BackendFeatures
+from spinta.backends.constants import BackendOrigin, BackendFeatures, DistributionType
 from spinta.core.ufuncs import Env
 from spinta.ufuncs.resultbuilder.components import ResultBuilder
 
@@ -58,3 +59,9 @@ class Backend:
 
 
 SelectTree = Optional[Dict[str, "SelectTree"]]
+
+
+@dataclasses.dataclass
+class DistributionStrategy:
+    distribution_type: DistributionType
+    property: str | None = None

spinta/backends/constants.py

@@ -35,3 +35,13 @@ class BackendFeatures(enum.Enum):
 
     # Backend supports
     EXPAND = "EXPAND"
+
+    # Backend supports sharding
+    DISTRIBUTE = "DISTRIBUTE"
+
+
+class DistributionType(enum.Enum):
+    SCHEMA = "schema"
+    TABLE = "table"
+    COPY = "copy"
+    UNDISTRIBUTED = "undistributed"

spinta/backends/helpers.py

@@ -26,7 +26,7 @@ from spinta.components import Model
 from spinta.components import Namespace
 from spinta.components import Property
 from spinta.exceptions import BackendUnavailable
-from spinta.types.datatype import DataType, Denorm
+from spinta.types.datatype import DataType, Denorm, String, Base32, PrimaryKey
 from spinta.utils.data import take
 from spinta.backends.constants import TableType, BackendOrigin
 
@@ -344,7 +344,7 @@ def get_ns_reserved_props(action: Action) -> list[str]:
     return []
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class TableIdentifier:
     """
     Represents a table identifier across logical (app) and PostgreSQL layers.
@@ -391,35 +391,42 @@ class TableIdentifier:
     table_arg: str | None = dataclasses.field(default=None)
     default_pg_schema: str | None = dataclasses.field(default=None)
 
-    logical_name: str = dataclasses.field(init=False)
+    logical_name: str = dataclasses.field(init=False, compare=False)
     # Name with namespace connected with '/', like it is used with Model class
-    logical_qualified_name: str = dataclasses.field(init=False)
+    logical_qualified_name: str = dataclasses.field(init=False, compare=False)
 
-    pg_table_name: str = dataclasses.field(init=False)
-    pg_schema_name: str | None = dataclasses.field(init=False)
+    pg_table_name: str = dataclasses.field(init=False, compare=False)
+    pg_schema_name: str | None = dataclasses.field(init=False, compare=False)
     # Used for hashed schema and table names
-    pg_qualified_name: str = dataclasses.field(init=False)
+    pg_qualified_name: str = dataclasses.field(init=False, compare=False)
     # Escaped qualified name, used for queries
-    pg_escaped_qualified_name: str = dataclasses.field(init=False)
+    pg_escaped_qualified_name: str = dataclasses.field(init=False, compare=False)
 
     def __post_init__(self):
-        self.logical_name = self.base_name + self.table_type.value
+        logical_name = self.base_name + self.table_type.value
         if self.table_arg:
-            self.logical_name += "/" + self.table_arg
+            logical_name += "/" + self.table_arg
 
-        self.logical_qualified_name = f"{self.schema}/{self.logical_name}" if self.schema else self.logical_name
+        logical_qualified_name = f"{self.schema}/{logical_name}" if self.schema else logical_name
 
-        self.pg_table_name = get_pg_name(self.logical_name)
-        self.pg_schema_name = get_pg_name(self.schema) if self.schema else self.default_pg_schema
-        self.pg_qualified_name = (
-            f"{self.pg_schema_name}.{self.pg_table_name}" if self.pg_schema_name else self.pg_table_name
-        )
-        self.pg_escaped_qualified_name = (
-            f"{pg_identifier_preparer.quote(self.pg_schema_name)}.{pg_identifier_preparer.quote(self.pg_table_name)}"
-            if self.pg_schema_name
-            else pg_identifier_preparer.quote(self.pg_table_name)
+        pg_table_name = get_pg_name(logical_name)
+        pg_schema_name = get_pg_name(self.schema) if self.schema else self.default_pg_schema
+        pg_qualified_name = f"{pg_schema_name}.{pg_table_name}" if pg_schema_name else pg_table_name
+        pg_escaped_qualified_name = (
+            f"{pg_identifier_preparer.quote(pg_schema_name)}.{pg_identifier_preparer.quote(pg_table_name)}"
+            if pg_schema_name
+            else pg_identifier_preparer.quote(pg_table_name)
         )
 
+        # This is needed because we want to make this dataclass hashable (frozen=True, does that)
+        # But because it becomes immutable, we need to set all the attributes manually (the same way dataclass __init__ does).
+        object.__setattr__(self, "logical_name", logical_name)
+        object.__setattr__(self, "logical_qualified_name", logical_qualified_name)
+        object.__setattr__(self, "pg_table_name", pg_table_name)
+        object.__setattr__(self, "pg_schema_name", pg_schema_name)
+        object.__setattr__(self, "pg_qualified_name", pg_qualified_name)
+        object.__setattr__(self, "pg_escaped_qualified_name", pg_escaped_qualified_name)
+
     def change_table_type(self, new_type: TableType, table_arg: str | None = None) -> "TableIdentifier":
         return dataclasses.replace(self, table_type=new_type, table_arg=table_arg)
 
@@ -596,3 +603,28 @@ def extract_table_data_from_logical_name(table_name: str) -> tuple[str | None, T
             return data[0], table_type, None
 
     return None, None, None
+
+
+def is_custom_id_prop(prop: Property) -> bool:
+    return prop.name == "_id" and not isinstance(prop.dtype, PrimaryKey)
+
+
+def is_custom_revision_prop(prop: Property) -> bool:
+    return prop.name == "_revision" and prop.explicitly_given
+
+
+def is_accessible_by_equals_sign(id_prop: Property, value: str | int) -> bool:
+    if isinstance(id_prop.dtype, Base32):
+        return True
+
+    if isinstance(id_prop.dtype, String):
+        return not check_if_model_primary_key_is_composite(id_prop.model)
+
+    return False
+
+
+def check_if_model_primary_key_is_composite(model: Model) -> bool:
+    pkeys_count = len(model.external.pkeys)
+    if pkeys_count > 1:
+        return True
+    return False

spinta/backends/postgresql/commands/wipe.py

@@ -48,7 +48,7 @@ def wipe(context: Context, model: Model, backend: PostgreSQL):
             if changelog_table_identifier.pg_schema_name
             else f'"{seqname}"'
         )
-        connection.execute(f"ALTER SEQUENCE {seq_escaped_named} RESTART")
+        connection.execute(sa.func.setval(seq_escaped_named, 1, False))
 
     # Delete data table
     table = backend.get_table(model)

spinta/backends/postgresql/components.py

@@ -26,7 +26,13 @@ class PostgreSQL(Backend):
         },
     }
 
-    features = {BackendFeatures.FILE_BLOCKS, BackendFeatures.WRITE, BackendFeatures.EXPAND, BackendFeatures.PAGINATION}
+    features = {
+        BackendFeatures.FILE_BLOCKS,
+        BackendFeatures.WRITE,
+        BackendFeatures.EXPAND,
+        BackendFeatures.PAGINATION,
+        BackendFeatures.DISTRIBUTE,
+    }
 
     engine: Engine = None
     schema: sa.MetaData = None

spinta/backends/postgresql/helpers/migrate/actions.py

@@ -3,7 +3,7 @@ from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.dialects import postgresql
 
 import sqlalchemy as sa
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Generator
 
 from spinta.backends.helpers import TableIdentifier
 from spinta.backends.postgresql.helpers.name import name_changed
@@ -91,6 +91,19 @@ class RenameTableMigrationAction(MigrationAction):
             )
 
 
+class SetTableCommentMigrationAction(MigrationAction):
+    def __init__(self, table_identifier: TableIdentifier, comment: str) -> None:
+        self.table_identifier = table_identifier
+        self.comment = comment
+
+    def execute(self, op: "Operations") -> None:
+        op.create_table_comment(
+            table_name=self.table_identifier.pg_table_name,
+            comment=self.comment,
+            schema=self.table_identifier.pg_schema_name,
+        )
+
+
 class AddColumnMigrationAction(MigrationAction):
     def __init__(self, table_identifier: TableIdentifier, column: sa.Column) -> None:
         self.table_identifier = table_identifier
@@ -149,6 +162,21 @@ class AlterColumnMigrationAction(MigrationAction):
         )
 
 
+class SetColumnCommentMigrationAction(MigrationAction):
+    def __init__(self, table_identifier: TableIdentifier, column: str, comment: str) -> None:
+        self.table_identifier = table_identifier
+        self.comment = comment
+        self.column = column
+
+    def execute(self, op: "Operations") -> None:
+        op.alter_column(
+            table_name=self.table_identifier.pg_table_name,
+            column_name=self.column,
+            comment=self.comment,
+            schema=self.table_identifier.pg_schema_name,
+        )
+
+
 class DropConstraintMigrationAction(MigrationAction):
     def __init__(self, table_identifier: TableIdentifier, constraint_name: str) -> None:
         self.table_identifier = table_identifier
@@ -547,6 +575,50 @@ class CreateSchemaMigrationAction(MigrationAction):
         op.execute(self.query)
 
 
+class DistributeSchema(MigrationAction):
+    def __init__(self, schema_name: str) -> None:
+        self.schema_name = schema_name
+        self.query = f"SELECT citus_schema_distribute('{pg_identifier_preparer.quote(schema_name)}')"
+
+    def execute(self, op: "Operations") -> None:
+        op.execute(self.query)
+
+
+class DistributeReference(MigrationAction):
+    def __init__(self, table_identifier: TableIdentifier) -> None:
+        self.query = f"SELECT create_reference_table('{table_identifier.pg_escaped_qualified_name}')"
+
+    def execute(self, op: "Operations") -> None:
+        op.execute(self.query)
+
+
+class DistributeTable(MigrationAction):
+    def __init__(self, table_identifier: TableIdentifier, column: str) -> None:
+        self.query = f"SELECT create_distributed_table('{table_identifier.pg_escaped_qualified_name}', '{column}')"
+
+    def execute(self, op: "Operations") -> None:
+        op.execute(self.query)
+
+
+class UndistributeSchema(MigrationAction):
+    def __init__(self, schema_name: str) -> None:
+        self.schema_name = schema_name
+        self.query = f"SELECT citus_schema_undistribute('{pg_identifier_preparer.quote(schema_name)}')"
+
+    def execute(self, op: "Operations") -> None:
+        op.execute(self.query)
+
+
+class UndistributeTable(MigrationAction):
+    def __init__(self, table_identifier: TableIdentifier) -> None:
+        self.query = (
+            f"SELECT undistribute_table('{table_identifier.pg_escaped_qualified_name}', cascade_via_foreign_keys=>true)"
+        )
+
+    def execute(self, op: "Operations") -> None:
+        op.execute(self.query)
+
+
 class MigrationHandler:
     def __init__(self) -> None:
         self.migrations: list[MigrationAction] = []
@@ -571,8 +643,15 @@ class MigrationHandler:
                     return True
         return False
 
-    def run_migrations(self, op: "Operations") -> None:
+    def gather_migrations(self) -> Generator[MigrationAction, None, None]:
         for migration in self.migrations:
-            migration.execute(op)
+            yield migration
         for migration in self.foreign_key_migration:
+            yield migration
+
+    def run_migrations(self, op: "Operations") -> None:
+        for migration in self.gather_migrations():
             migration.execute(op)
+
+    def count(self) -> int:
+        return len(list(self.gather_migrations()))