use core::fmt;
use std::{borrow::Cow, collections::BTreeSet, marker::PhantomData};
use super::{
ApplyExpressionNames, ApplyProjection, AttrNames, Expression, ProjectionableBuilder,
fmt_attr_maps, resolve_expression, utils::resolve_attr_path,
};
use crate::{
AttributeDefinition, CompositeKey, CompositeKeySchema, KeySchema, KeySchemaKind, SimpleKey,
SimpleKeySchema, TableDefinition,
};
#[derive(Debug, Clone)]
struct BuiltProjection {
expression: Expression,
names: AttrNames,
}
#[derive(Debug, Clone)]
#[must_use = "expression does nothing until applied to a request"]
pub struct Projection<'a, TD> {
attrs: BTreeSet<Cow<'a, str>>,
_marker: PhantomData<TD>,
}
impl<'a, TD: TableDefinition> Projection<'a, TD>
where
Self: key_schema_projection::KeySchemaProjection<
'a,
TD::KeySchema,
<TD::KeySchema as KeySchema>::Kind,
>,
{
pub fn new(attrs: impl IntoIterator<Item = impl Into<Cow<'a, str>>>) -> Self {
Self {
attrs: <Self as key_schema_projection::KeySchemaProjection<
'a,
TD::KeySchema,
<TD::KeySchema as KeySchema>::Kind,
>>::key_schema_names()
.chain(attrs.into_iter().map(Into::into))
.collect(),
_marker: PhantomData,
}
}
pub fn keys_only() -> Self {
Self::new([] as [Cow<'a, str>; 0])
}
}
mod key_schema_projection {
use super::*;
pub trait KeySchemaProjection<'a, KS: KeySchema, KSK: KeySchemaKind> {
fn key_schema_names() -> impl Iterator<Item = Cow<'a, str>>;
}
impl<'a, TD: TableDefinition> KeySchemaProjection<'a, TD::KeySchema, SimpleKey>
for Projection<'a, TD>
where
TD::KeySchema: SimpleKeySchema,
{
fn key_schema_names() -> impl Iterator<Item = Cow<'a, str>> {
[<<TD::KeySchema as KeySchema>::PartitionKey as AttributeDefinition>::NAME.into()]
.into_iter()
}
}
impl<'a, TD: TableDefinition> KeySchemaProjection<'a, TD::KeySchema, CompositeKey>
for Projection<'a, TD>
where
TD::KeySchema: CompositeKeySchema,
{
fn key_schema_names() -> impl Iterator<Item = Cow<'a, str>> {
[
<<TD::KeySchema as KeySchema>::PartitionKey as AttributeDefinition>::NAME.into(),
<<TD::KeySchema as CompositeKeySchema>::SortKey as AttributeDefinition>::NAME
.into(),
]
.into_iter()
}
}
}
impl<TD> Projection<'_, TD> {
fn build(&self) -> BuiltProjection {
let mut counter = 0;
let mut all_names = Vec::new();
let mut resolved_parts = Vec::new();
for attr in &self.attrs {
let (expr, names) = resolve_attr_path(attr, "p", &mut counter);
all_names.extend(names);
resolved_parts.push(expr.into_owned());
}
BuiltProjection {
expression: resolved_parts.join(","),
names: all_names,
}
}
}
impl<TD> fmt::Display for Projection<'_, TD> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let built = self.build();
if built.expression.is_empty() {
return f.write_str("<none>");
}
let no_values = vec![];
if f.alternate() {
f.write_str(&built.expression)?;
fmt_attr_maps(f, &built.names, &no_values)
} else {
f.write_str(&resolve_expression(
&built.expression,
&built.names,
&no_values,
))
}
}
}
impl<TD, B: ProjectionableBuilder> ApplyProjection<B> for Projection<'_, TD> {
fn apply_projection(self, builder: B) -> B {
let built = self.build();
if built.expression.is_empty() {
return builder;
}
builder
.projection_expression(built.expression)
.apply_names(built.names)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_fixtures::*;
crate::attribute_definitions! {
MyPk { "MyPk": crate::StringAttribute }
}
crate::table_definitions! {
SimpleTable {
type PartitionKey = MyPk;
fn table_name() -> String { "simple".to_owned() }
}
}
#[test]
fn test_projection_new_simple_key_auto_includes_pk() {
let proj = Projection::<SimpleTable>::new(["name", "email"]);
let output = format!("{proj}");
assert_eq!(output, "MyPk,email,name");
}
#[test]
fn test_projection_new_composite_key_auto_includes_pk_and_sk() {
let proj = Projection::<PlatformTable>::new(["email", "role"]);
let output = format!("{proj}");
assert_eq!(output, "PK,SK,email,role");
}
#[test]
fn test_projection_new_dedup_when_user_supplies_pk() {
let proj = Projection::<PlatformTable>::new(["PK", "SK", "email"]);
let output = format!("{proj}");
assert_eq!(output, "PK,SK,email");
}
#[test]
fn test_projection_display_default_with_reserved_word() {
let proj = Projection::<PlatformTable>::new(["email", "Status"]);
let output = format!("{proj}");
assert_eq!(output, "PK,SK,Status,email");
}
#[test]
fn test_projection_display_alternate_with_reserved_word() {
let proj = Projection::<PlatformTable>::new(["email", "Status"]);
let output = format!("{proj:#}");
assert_eq!(output, "PK,SK,#p0,email\n names: { #p0 = Status }");
}
}