spinta 0.2.dev24__tar.gz → 0.2.dev26__tar.gz

{spinta-0.2.dev24 → spinta-0.2.dev26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spinta
-Version: 0.2.dev24
+Version: 0.2.dev26
 Summary: A platform for describing, extracting, transforming, loading and serving open data.
 License: MIT
 License-File: LICENSE

{spinta-0.2.dev24 → spinta-0.2.dev26}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "spinta"
-version = "0.2dev24"
+version = "0.2dev26"
 description = "A platform for describing, extracting, transforming, loading and serving open data."
 authors = ["Mantas Zimnickas <sirexas@gmail.com>"]
 license = "MIT"

{spinta-0.2.dev24 → spinta-0.2.dev26}/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-0.2.dev24 → spinta-0.2.dev26}/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-0.2.dev24 → spinta-0.2.dev26}/spinta/backends/helpers.py

@@ -91,7 +91,7 @@ def get_select_tree(
         select = list(select.keys())
 
     select = _apply_always_show_id(context, action, select)
-    if select is None and action in (Action.GETALL, Action.SEARCH):
+    if select is None and action in (Action.GETALL, Action.SEARCH, Action.GETONE):
         # If select is not given, select everything.
         select = {"*": {}}
     return flat_select_to_nested(select)
@@ -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)
 

{spinta-0.2.dev24 → spinta-0.2.dev26}/spinta/backends/postgresql/commands/load.py

@@ -5,6 +5,7 @@ import sqlalchemy as sa
 from spinta import commands
 from spinta.backends.postgresql.components import PostgreSQL
 from spinta.backends.postgresql.helpers.name import PG_NAMING_CONVENTION
+from spinta.backends.postgresql.sqlalchemy import create_postgresql_engine
 from spinta.components import Context
 from spinta.utils.sqlalchemy import get_metadata_naming_convention
 
@@ -12,7 +13,7 @@ from spinta.utils.sqlalchemy import get_metadata_naming_convention
 @commands.load.register(Context, PostgreSQL, dict)
 def load(context: Context, backend: PostgreSQL, config: Dict[str, Any]):
     backend.dsn = config["dsn"]
-    backend.engine = sa.create_engine(backend.dsn, echo=False)
+    backend.engine = create_postgresql_engine(backend.dsn, echo=False)
     backend.schema = sa.MetaData(backend.engine, naming_convention=get_metadata_naming_convention(PG_NAMING_CONVENTION))
     backend.tables = {}
 

{spinta-0.2.dev24 → spinta-0.2.dev26}/spinta/backends/postgresql/commands/wait.py

@@ -1,7 +1,7 @@
-import sqlalchemy as sa
 import sqlalchemy.exc
 
 from spinta import commands
+from spinta.backends.postgresql.sqlalchemy import create_postgresql_engine
 from spinta.components import Context
 from spinta.backends.postgresql.components import PostgreSQL
 
@@ -10,7 +10,7 @@ from spinta.backends.postgresql.components import PostgreSQL
 def wait(context: Context, backend: PostgreSQL, *, fail: bool = False) -> bool:
     rc = context.get("rc")
     dsn = rc.get("backends", backend.name, "dsn", required=True)
-    engine = sa.create_engine(dsn, connect_args={"connect_timeout": 0})
+    engine = create_postgresql_engine(dsn, connect_args={"connect_timeout": 0})
     try:
         conn = engine.connect()
     except sqlalchemy.exc.OperationalError:

{spinta-0.2.dev24 → spinta-0.2.dev26}/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-0.2.dev24 → spinta-0.2.dev26}/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-0.2.dev24 → spinta-0.2.dev26}/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()))