Skip to main content

canic_sharding_runtime/api/
mod.rs

1use crate::workflow::{ShardingWorkflow, query::ShardingQuery};
2use canic_core::{
3    cdk::types::Principal,
4    dto::{
5        error::Error,
6        placement::sharding::{
7            ShardingPartitionKeysResponse, ShardingPlanStateResponse, ShardingRegistryResponse,
8        },
9    },
10};
11
12///
13/// ShardingApi
14///
15
16pub struct ShardingApi;
17
18impl ShardingApi {
19    /// Lookup the shard assigned to a partition_key in a pool, if any.
20    #[must_use]
21    pub fn lookup_partition_key(pool: &str, partition_key: &str) -> Option<Principal> {
22        ShardingQuery::lookup_partition_key(pool, partition_key)
23    }
24
25    /// Return the shard for a partition_key, or an Error if unassigned.
26    pub fn resolve_shard_for_key(
27        pool: &str,
28        partition_key: impl AsRef<str>,
29    ) -> Result<Principal, Error> {
30        ShardingQuery::resolve_shard_for_key(pool, partition_key.as_ref()).map_err(Error::from)
31    }
32
33    /// Return a view of the full sharding registry.
34    #[must_use]
35    pub fn registry() -> ShardingRegistryResponse {
36        ShardingQuery::registry()
37    }
38
39    /// Return all partition_keys currently assigned to a shard.
40    #[must_use]
41    pub fn partition_keys(pool: &str, shard: Principal) -> ShardingPartitionKeysResponse {
42        ShardingQuery::partition_keys(pool, shard)
43    }
44
45    /// Assign a partition_key to a shard in the given pool.
46    pub async fn assign_to_pool(
47        pool: &str,
48        partition_key: impl AsRef<str>,
49    ) -> Result<Principal, Error> {
50        ShardingWorkflow::assign_to_pool(pool, partition_key)
51            .await
52            .map_err(Error::from)
53    }
54
55    /// Perform a dry-run shard assignment and return the resulting plan.
56    pub fn plan_assign_to_pool(
57        pool: &str,
58        partition_key: impl AsRef<str>,
59    ) -> Result<ShardingPlanStateResponse, Error> {
60        ShardingWorkflow::plan_assign_to_pool(pool, partition_key).map_err(Error::from)
61    }
62}