spinta 0.2.dev24__py3-none-any.whl → 0.2.dev26__py3-none-any.whl

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

@@ -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/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/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/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()))

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

@@ -0,0 +1,383 @@
+import dataclasses
+from collections import defaultdict
+from copy import deepcopy
+
+import sqlalchemy as sa
+from tqdm import tqdm
+from multipledispatch import dispatch
+
+from spinta.backends import Backend
+from spinta.backends.constants import DistributionType
+from spinta.backends.helpers import TableIdentifier
+from spinta.backends.helpers import get_table_identifier
+from spinta.backends.postgresql.components import PostgreSQL
+from spinta.backends.postgresql.helpers.migrate.actions import (
+    MigrationHandler,
+    UndistributeSchema,
+    UndistributeTable,
+    DistributeReference,
+    DistributeTable,
+    DistributeSchema,
+)
+from spinta.cli.helpers.message import cli_message
+from spinta.components import Context, Model
+from spinta.exceptions import NotImplementedFeature
+
+
+@dataclasses.dataclass
+class ShardingPlan:
+    schemas: set[str] = dataclasses.field(default_factory=set)
+    references: set[TableIdentifier] = dataclasses.field(default_factory=set)
+    distributed: dict[TableIdentifier, str] = dataclasses.field(default_factory=dict)
+    local: set[TableIdentifier] = dataclasses.field(default_factory=set)
+
+    _lookup: dict[TableIdentifier | str, DistributionType] = dataclasses.field(init=False, default_factory=dict)
+
+    def __sub__(self, other) -> "ShardingPlan":
+        return ShardingPlan(
+            schemas=self.schemas - other.schemas,
+            references=self.references - other.references,
+            distributed=dict(self.distributed.items() - other.distributed.items()),
+            local=self.local - other.local,
+        )
+
+    def __post_init__(self) -> None:
+        for schema in self.schemas:
+            self._lookup[schema] = DistributionType.SCHEMA
+
+        for table_identifier in self.distributed.keys():
+            self._lookup[table_identifier] = DistributionType.TABLE
+
+        for table_identifier in self.references:
+            self._lookup[table_identifier] = DistributionType.COPY
+
+        for table_identifier in self.local:
+            self._lookup[table_identifier] = DistributionType.UNDISTRIBUTED
+
+    def distribution_type(self, key: TableIdentifier | str) -> DistributionType | None:
+        return self._lookup.get(key, None)
+
+    def discard(self, key: TableIdentifier | str) -> None:
+        distribution_type = self._lookup.pop(key, None)
+        if distribution_type is None:
+            return
+
+        match distribution_type:
+            case DistributionType.SCHEMA:
+                self.schemas.discard(key)
+            case DistributionType.TABLE:
+                self.distributed.pop(key)
+            case DistributionType.COPY:
+                self.references.discard(key)
+            case DistributionType.UNDISTRIBUTED:
+                self.local.discard(key)
+
+
+def _generate_current_distribution_query(schemas: list[str] | None = None) -> (str, dict):
+    base_query = """
+    SELECT
+        n.nspname AS schema_name,
+        c.relname AS table_name,
+        format('%I.%I', n.nspname, c.relname) AS full_table_name,
+        d.description AS table_comment,
+        COALESCE(ct.citus_table_type, 'local') AS distribution_type,
+        ct.distribution_column
+    FROM pg_class c
+    JOIN pg_namespace n ON n.oid = c.relnamespace
+    LEFT JOIN citus_tables ct ON ct.table_name = c.oid::regclass
+    LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = 0
+    WHERE c.relkind = 'r'
+      AND n.nspname NOT LIKE 'pg_%'
+      AND n.nspname <> 'information_schema'
+    """
+
+    params = {}
+    if schemas:
+        base_query += " AND n.nspname = ANY(:schemas)"
+        params["schemas"] = schemas
+
+    base_query += " ORDER BY n.nspname, c.relname"
+    return base_query, params
+
+
+@dispatch(Context)
+def gather_current_sharding_plan(context: Context, **kwargs) -> dict[str, ShardingPlan]:
+    store = context.get("store")
+    backends = store.backends
+
+    plans: dict[str, ShardingPlan] = defaultdict(ShardingPlan)
+    for backend_name, backend in backends.items():
+        plans[backend_name] = gather_current_sharding_plan(context, backend, **kwargs)
+    return plans
+
+
+@dispatch(Context, Backend)
+def gather_current_sharding_plan(context: Context, backend: Backend, **kwargs) -> ShardingPlan:
+    # Currently, only postgresql backend supports citus distribution, instead of erroring, return empty plan.
+    return ShardingPlan()
+
+
+@dispatch(Context, PostgreSQL)
+def gather_current_sharding_plan(
+    context: Context, backend: PostgreSQL, schemas: list[str] | None = None, **kwargs
+) -> ShardingPlan:
+    plan = ShardingPlan()
+    with backend.begin() as conn:
+        query, params = _generate_current_distribution_query(schemas)
+        rows = conn.execute(sa.text(query), params).fetchall()
+
+        for row in rows:
+            if not row["table_comment"]:
+                continue
+
+            table_identifier = get_table_identifier(row["table_comment"])
+            match row["distribution_type"]:
+                case "schema":
+                    plan.schemas.add(row["schema_name"])
+                case "distributed":
+                    plan.distributed[table_identifier] = row["distribution_column"]
+                case "reference":
+                    plan.references.add(table_identifier)
+                case _:
+                    plan.local.add(table_identifier)
+    return plan
+
+
+def create_sharding_plan(context: Context, models: list[Model], **kwargs) -> dict[str, ShardingPlan]:
+    plans = defaultdict(ShardingPlan)
+    for model in models:
+        if not isinstance(model.backend, PostgreSQL):
+            continue
+
+        if not model.external or not model.external.dataset:
+            continue
+
+        distribution_strategy = model.distribution_strategy
+        plan = plans[model.backend.name]
+        table_identifier = get_table_identifier(model)
+        match distribution_strategy.distribution_type:
+            case DistributionType.SCHEMA:
+                plan.schemas.add(table_identifier.pg_schema_name)
+            case DistributionType.TABLE:
+                plan.distributed[table_identifier] = distribution_strategy.property
+            case DistributionType.COPY:
+                plan.references.add(table_identifier)
+            case _:
+                plan.local.add(table_identifier)
+
+    return plans
+
+
+def invalidate_default_distribution(
+    context: Context, backend: PostgreSQL, plan: ShardingPlan, verbose: bool = False, **kwargs
+) -> ShardingPlan:
+    default_distribution = context.get("config").default_distribution_strategy
+    if not default_distribution:
+        return plan
+
+    match default_distribution.distribution_type:
+        case DistributionType.SCHEMA:
+            return invalidate_default_schema_distributions(context, backend, plan)
+        case _:
+            if verbose:
+                cli_message(
+                    f"Skipped invalidation of default distribution for {default_distribution.distribution_type.value} type"
+                )
+
+            return plan
+
+
+def valid_schema_distribution_foreign_key(plan: ShardingPlan, schema: str, foreign_key: dict) -> bool:
+    if foreign_key["referred_schema"] == schema:
+        return True
+
+    for reference in plan.references:
+        if (
+            reference.pg_schema_name == foreign_key["referred_schema"]
+            and reference.pg_table_name == foreign_key["referred_table"]
+        ):
+            return True
+
+    return False
+
+
+@dispatch(Context, Backend, ShardingPlan)
+def invalidate_default_schema_distributions(
+    context: Context, backend: Backend, plan: ShardingPlan, **kwargs
+) -> ShardingPlan:
+    raise NotImplementedFeature(f"Ability to invalidate default schema distribution for {backend.type!r} backend type")
+
+
+@dispatch(Context, PostgreSQL, ShardingPlan)
+def invalidate_default_schema_distributions(
+    context: Context, backend: PostgreSQL, plan: ShardingPlan, **kwargs
+) -> ShardingPlan:
+    if not plan.schemas:
+        return plan
+
+    invalid_schemas = set()
+
+    inspector = sa.inspect(backend.engine)
+
+    plan_copy = deepcopy(plan)
+    for schema in plan.schemas:
+        tables = inspector.get_table_names(schema=schema)
+
+        for table in tables:
+            foreign_keys = inspector.get_foreign_keys(table, schema=schema)
+            if not foreign_keys:
+                continue
+
+            for key in foreign_keys:
+                if not valid_schema_distribution_foreign_key(plan, schema, key):
+                    invalid_schemas.add(key["referred_schema"])
+                    invalid_schemas.add(schema)
+
+    if not invalid_schemas:
+        return plan
+
+    for invalid_schema in invalid_schemas:
+        plan_copy.schemas.discard(invalid_schema)
+
+    return plan_copy
+
+
+def _build_fk_graph(conn: sa.engine.Connection) -> dict[tuple[str, str], set[tuple[str, str]]]:
+    graph: dict[tuple[str, str], set[tuple[str, str]]] = defaultdict(set)
+
+    rows = conn.execute("""
+                        SELECT
+                            src_ns.nspname AS source_schema,
+                            src.relname    AS source_table,
+                            tgt_ns.nspname AS target_schema,
+                            tgt.relname    AS target_table
+                        FROM pg_constraint con
+                                 JOIN pg_class src
+                                      ON src.oid = con.conrelid
+                                 JOIN pg_namespace src_ns
+                                      ON src_ns.oid = src.relnamespace
+                                 JOIN pg_class tgt
+                                      ON tgt.oid = con.confrelid
+                                 JOIN pg_namespace tgt_ns
+                                      ON tgt_ns.oid = tgt.relnamespace
+                        WHERE con.contype = 'f'
+                        """)
+
+    for src_schema, src_table, tgt_schema, tgt_table in rows.fetchall():
+        src = (src_schema, src_table)
+        tgt = (tgt_schema, tgt_table)
+
+        if src == tgt:
+            continue
+
+        graph[src].add(tgt)
+        graph[tgt].add(src)
+
+    visited: set[tuple[str, str]] = set()
+    key_to_component: dict[tuple[str, str], set[tuple[str, str]]] = {}
+
+    for node in graph:
+        if node in visited:
+            continue
+
+        stack = [node]
+        component: set[tuple[str, str]] = set()
+
+        while stack:
+            cur = stack.pop()
+            if cur in visited:
+                continue
+
+            visited.add(cur)
+            component.add(cur)
+            stack.extend(graph[cur] - visited)
+
+        for n in component:
+            key_to_component[n] = component
+
+    return key_to_component
+
+
+def build_fk_components(
+    conn: sa.engine.Connection, tables: set[TableIdentifier]
+) -> dict[TableIdentifier, set[TableIdentifier]]:
+    """
+    Build FK-connected components using FULL DB graph,
+    then return mapping only for given `tables`.
+    """
+
+    target_map = {(t.pg_schema_name, t.pg_table_name): t for t in tables}
+    graph = _build_fk_graph(conn)
+
+    result = {}
+    for key, table in target_map.items():
+        full_component = graph.get(key, {key})
+        result[table] = {target_map[k] for k in full_component if k in target_map}
+
+    return result
+
+
+def undistribute_all(
+    context: Context,
+    backend: PostgreSQL,
+    plan: ShardingPlan,
+    handler: MigrationHandler,
+    progress_bar: tqdm | None = None,
+    **kwargs,
+) -> None:
+    for schema in plan.schemas:
+        handler.add_action(UndistributeSchema(schema_name=schema))
+        if progress_bar is not None:
+            progress_bar.update(1)
+
+    if not (plan.references or plan.distributed):
+        return
+
+    processed = set()
+    undistributed_tables = plan.references | set(table for table in plan.distributed.keys())
+    with backend.begin() as conn:
+        component_map = build_fk_components(conn, undistributed_tables)
+
+    for table in plan.distributed.keys():
+        if table in processed:
+            continue
+
+        handler.add_action(UndistributeTable(table_identifier=table))
+        if progress_bar is not None:
+            progress_bar.update(1)
+        component = component_map[table]
+        processed.update(component)
+
+    for table in plan.references:
+        if table in processed:
+            continue
+
+        handler.add_action(UndistributeTable(table_identifier=table))
+        if progress_bar is not None:
+            progress_bar.update(1)
+        component = component_map[table]
+        processed.update(component)
+
+
+def distribute_all(
+    context: Context,
+    backend: PostgreSQL,
+    plan: ShardingPlan,
+    handler: MigrationHandler,
+    progress_bar: tqdm | None = None,
+    **kwargs,
+) -> None:
+    for table in plan.references:
+        handler.add_action(DistributeReference(table_identifier=table))
+        if progress_bar is not None:
+            progress_bar.update(1)
+
+    for table, column in plan.distributed.items():
+        handler.add_action(DistributeTable(table_identifier=table, column=column))
+        if progress_bar is not None:
+            progress_bar.update(1)
+
+    for schema in plan.schemas:
+        handler.add_action(DistributeSchema(schema_name=schema))
+        if progress_bar is not None:
+            progress_bar.update(1)