Skip to main content

use_partition_key/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../README.md")]
3
4use core::fmt;
5
6macro_rules! string_newtype {
7    ($(#[$meta:meta])* $name:ident) => {
8        $(#[$meta])*
9        #[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
10        pub struct $name(String);
11
12        impl $name {
13            /// Creates a new string-backed primitive.
14            pub fn new(value: impl Into<String>) -> Self {
15                Self(value.into())
16            }
17
18            /// Returns the stored string value.
19            pub fn as_str(&self) -> &str {
20                &self.0
21            }
22        }
23
24        impl AsRef<str> for $name {
25            fn as_ref(&self) -> &str {
26                self.as_str()
27            }
28        }
29
30        impl From<String> for $name {
31            fn from(value: String) -> Self {
32                Self::new(value)
33            }
34        }
35
36        impl From<&str> for $name {
37            fn from(value: &str) -> Self {
38                Self::new(value)
39            }
40        }
41
42        impl fmt::Display for $name {
43            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
44                formatter.write_str(self.as_str())
45            }
46        }
47    };
48}
49
50string_newtype! {
51    /// A reusable partition key.
52    PartitionKey
53}
54string_newtype! {
55    /// A shard key.
56    ShardKey
57}
58string_newtype! {
59    /// A routing key.
60    RoutingKey
61}
62string_newtype! {
63    /// A sort key.
64    SortKey
65}
66
67/// A composite partitioning key.
68#[derive(Clone, Debug, Default, Eq, PartialEq)]
69pub struct CompositeKey {
70    parts: Vec<PartitionKey>,
71}
72
73impl CompositeKey {
74    /// Creates an empty composite key.
75    pub fn new() -> Self {
76        Self { parts: Vec::new() }
77    }
78
79    /// Creates a composite key from parts.
80    pub fn from_parts(parts: Vec<PartitionKey>) -> Self {
81        Self { parts }
82    }
83
84    /// Adds a partition key part.
85    pub fn with_part(mut self, part: PartitionKey) -> Self {
86        self.parts.push(part);
87        self
88    }
89
90    /// Returns all parts.
91    pub fn parts(&self) -> &[PartitionKey] {
92        &self.parts
93    }
94}
95
96/// Partition strategy labels.
97#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
98pub enum PartitionStrategy {
99    Hash,
100    Range,
101    List,
102    Composite,
103    Manual,
104    #[default]
105    Unknown,
106}
107
108impl PartitionStrategy {
109    /// Returns a stable lowercase label.
110    pub const fn as_str(self) -> &'static str {
111        match self {
112            Self::Hash => "hash",
113            Self::Range => "range",
114            Self::List => "list",
115            Self::Composite => "composite",
116            Self::Manual => "manual",
117            Self::Unknown => "unknown",
118        }
119    }
120}
121
122impl fmt::Display for PartitionStrategy {
123    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
124        formatter.write_str(self.as_str())
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::{CompositeKey, PartitionKey, PartitionStrategy, RoutingKey, ShardKey, SortKey};
131    use std::collections::hash_map::DefaultHasher;
132    use std::hash::{Hash, Hasher};
133
134    #[test]
135    fn constructs_partition_keys() {
136        assert_eq!(PartitionKey::new("tenant_1").to_string(), "tenant_1");
137        assert_eq!(ShardKey::new("shard-a").as_ref(), "shard-a");
138        assert_eq!(RoutingKey::new("route-1").as_str(), "route-1");
139        assert_eq!(SortKey::new("created_at").to_string(), "created_at");
140    }
141
142    #[test]
143    fn builds_composite_keys_and_formats_strategies() {
144        let key = CompositeKey::new()
145            .with_part(PartitionKey::new("tenant_1"))
146            .with_part(PartitionKey::new("customer_1"));
147
148        assert_eq!(key.parts().len(), 2);
149        assert_eq!(PartitionStrategy::Composite.to_string(), "composite");
150    }
151
152    #[test]
153    fn hashes_equal_partition_keys() {
154        let mut left = DefaultHasher::new();
155        let mut right = DefaultHasher::new();
156        PartitionKey::new("same").hash(&mut left);
157        PartitionKey::new("same").hash(&mut right);
158        assert_eq!(left.finish(), right.finish());
159    }
160}