use new IntoSelector trait to identify keys

This commit is contained in:
TotallyNot 2023-06-11 15:29:15 +02:00
parent 3e3909bb6c
commit 0799d6d475
3 changed files with 313 additions and 186 deletions

View file

@ -30,7 +30,7 @@ where
}
pub trait ApiKey: Sync + Send {
type IdType: PartialEq + Eq + std::hash::Hash;
type IdType: PartialEq + Eq + std::hash::Hash + Send + Sync;
fn value(&self) -> &str;
@ -44,12 +44,65 @@ pub trait KeyDomain: Clone + std::fmt::Debug + Send + Sync {
}
#[derive(Debug, Clone)]
pub enum KeySelector<K>
pub enum KeySelector<K, D>
where
K: ApiKey,
D: KeyDomain,
{
Key(String),
Id(K::IdType),
UserId(i32),
Has(D),
OneOf(Vec<D>),
}
impl<K, D> KeySelector<K, D>
where
K: ApiKey,
D: KeyDomain,
{
pub(crate) fn fallback(&self) -> Option<Self> {
match self {
Self::Key(_) | Self::UserId(_) | Self::Id(_) => None,
Self::Has(domain) => domain.fallback().map(Self::Has),
Self::OneOf(domains) => {
let fallbacks: Vec<_> = domains.into_iter().filter_map(|d| d.fallback()).collect();
if fallbacks.is_empty() {
None
} else {
Some(Self::OneOf(fallbacks))
}
}
}
}
}
pub trait IntoSelector<K, D>: Send + Sync
where
K: ApiKey,
D: KeyDomain,
{
fn into_selector(self) -> KeySelector<K, D>;
}
impl<K, D> IntoSelector<K, D> for D
where
K: ApiKey,
D: KeyDomain,
{
fn into_selector(self) -> KeySelector<K, D> {
KeySelector::Has(self)
}
}
impl<K, D> IntoSelector<K, D> for KeySelector<K, D>
where
K: ApiKey,
D: KeyDomain,
{
fn into_selector(self) -> KeySelector<K, D> {
self
}
}
#[async_trait]
@ -58,13 +111,17 @@ pub trait KeyPoolStorage {
type Domain: KeyDomain;
type Error: std::error::Error + Sync + Send;
async fn acquire_key(&self, domain: Self::Domain) -> Result<Self::Key, Self::Error>;
async fn acquire_key<S>(&self, selector: S) -> Result<Self::Key, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn acquire_many_keys(
async fn acquire_many_keys<S>(
&self,
domain: Self::Domain,
selector: S,
number: i64,
) -> Result<Vec<Self::Key>, Self::Error>;
) -> Result<Vec<Self::Key>, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn flag_key(&self, key: Self::Key, code: u8) -> Result<bool, Self::Error>;
@ -75,34 +132,41 @@ pub trait KeyPoolStorage {
domains: Vec<Self::Domain>,
) -> Result<Self::Key, Self::Error>;
async fn read_key(&self, key: KeySelector<Self::Key>)
-> Result<Option<Self::Key>, Self::Error>;
async fn read_key<S>(&self, selector: S) -> Result<Option<Self::Key>, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn read_user_keys(&self, user_id: i32) -> Result<Vec<Self::Key>, Self::Error>;
async fn read_keys<S>(&self, selector: S) -> Result<Vec<Self::Key>, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn remove_key(&self, key: KeySelector<Self::Key>) -> Result<Self::Key, Self::Error>;
async fn remove_key<S>(&self, selector: S) -> Result<Self::Key, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn query_key(&self, domain: Self::Domain) -> Result<Option<Self::Key>, Self::Error>;
async fn query_all(&self, domain: Self::Domain) -> Result<Vec<Self::Key>, Self::Error>;
async fn add_domain_to_key(
async fn add_domain_to_key<S>(
&self,
key: KeySelector<Self::Key>,
selector: S,
domain: Self::Domain,
) -> Result<Self::Key, Self::Error>;
) -> Result<Self::Key, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn remove_domain_from_key(
async fn remove_domain_from_key<S>(
&self,
key: KeySelector<Self::Key>,
selector: S,
domain: Self::Domain,
) -> Result<Self::Key, Self::Error>;
) -> Result<Self::Key, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
async fn set_domains_for_key(
async fn set_domains_for_key<S>(
&self,
key: KeySelector<Self::Key>,
selector: S,
domains: Vec<Self::Domain>,
) -> Result<Self::Key, Self::Error>;
) -> Result<Self::Key, Self::Error>
where
S: IntoSelector<Self::Key, Self::Domain>;
}
#[derive(Debug, Clone)]