Skip to content

Commit

Permalink
chore: impl Display for LicenseInfo and JWTClaims<T> (#16244)
Browse files Browse the repository at this point in the history
  • Loading branch information
drmingdrmer authored Aug 14, 2024
1 parent 4db95ee commit 66c8dea
Show file tree
Hide file tree
Showing 6 changed files with 290 additions and 24 deletions.
95 changes: 95 additions & 0 deletions src/common/base/src/display/display_slice.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt;

/// Implement `Display` for `&[T]` if T is `Display`.
///
/// It outputs at most `MAX` elements, excluding those from the 5th to the second-to-last one:
/// - `DisplaySlice(&[1,2,3,4,5,6])` outputs: `"[1,2,3,4,...,6]"`.
pub struct DisplaySlice<'a, T: fmt::Display, const MAX: usize = 5>(pub &'a [T]);

impl<'a, T: fmt::Display, const MAX: usize> fmt::Display for DisplaySlice<'a, T, MAX> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let slice = self.0;
let len = slice.len();

write!(f, "[")?;

if len > MAX {
for (i, t) in slice[..(MAX - 1)].iter().enumerate() {
if i > 0 {
write!(f, ",")?;
}

write!(f, "{}", t)?;
}

write!(f, ",..,")?;
write!(f, "{}", slice.last().unwrap())?;
} else {
for (i, t) in slice.iter().enumerate() {
if i > 0 {
write!(f, ",")?;
}

write!(f, "{}", t)?;
}
}

write!(f, "]")
}
}

pub trait DisplaySliceExt<'a, T: fmt::Display> {
fn display(&'a self) -> DisplaySlice<'a, T>;

/// Display at most `MAX` elements.
fn display_n<const MAX: usize>(&'a self) -> DisplaySlice<'a, T, MAX>;
}

impl<T> DisplaySliceExt<'_, T> for [T]
where T: fmt::Display
{
fn display(&self) -> DisplaySlice<T> {
DisplaySlice(self)
}

fn display_n<const MAX: usize>(&'_ self) -> DisplaySlice<'_, T, MAX> {
DisplaySlice(self)
}
}

#[cfg(test)]
mod tests {
use super::DisplaySlice;

#[test]
fn test_display_slice() {
let a = vec![1, 2, 3, 4];
assert_eq!("[1,2,3,4]", DisplaySlice::<_>(&a).to_string());

let a = vec![1, 2, 3, 4, 5];
assert_eq!("[1,2,3,4,5]", DisplaySlice::<_>(&a).to_string());

let a = vec![1, 2, 3, 4, 5, 6];
assert_eq!("[1,2,3,4,..,6]", DisplaySlice::<_>(&a).to_string());

let a = vec![1, 2, 3, 4, 5, 6, 7];
assert_eq!("[1,2,3,4,..,7]", DisplaySlice::<_>(&a).to_string());

let a = vec![1, 2, 3, 4, 5, 6, 7];
assert_eq!("[1,..,7]", DisplaySlice::<_, 2>(&a).to_string());
}
}
1 change: 1 addition & 0 deletions src/common/base/src/display/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,5 @@
// limitations under the License.

pub mod display_option;
pub mod display_slice;
pub mod display_unix_epoch;
94 changes: 94 additions & 0 deletions src/common/license/src/display_jwt_claims.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt;
use std::time::Duration;

use databend_common_base::display::display_option::DisplayOptionExt;
use databend_common_base::display::display_unix_epoch::DisplayUnixTimeStamp;
use databend_common_base::display::display_unix_epoch::DisplayUnixTimeStampExt;
use jwt_simple::claims::JWTClaims;
use jwt_simple::prelude::coarsetime;

/// A helper struct to Display JWT claims.
pub struct DisplayJWTClaims<'a, T> {
claims: &'a JWTClaims<T>,
}

impl<'a, T> fmt::Display for DisplayJWTClaims<'a, T>
where T: fmt::Display
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fn to_ts(v: Option<coarsetime::UnixTimeStamp>) -> Option<DisplayUnixTimeStamp> {
v.map(|x| {
let duration = Duration::from_micros(x.as_micros());
duration.display_unix_timestamp()
})
}

write!(f, "JWTClaims{{")?;
write!(f, "issuer: {}, ", self.claims.issuer.display())?;
write!(f, "issued_at: {}, ", to_ts(self.claims.issued_at).display())?;
write!(
f,
"expires_at: {}, ",
to_ts(self.claims.expires_at).display()
)?;
write!(f, "custom: {}", self.claims.custom)?;
write!(f, "}}")
}
}

/// Add `display_jwt_claims` method to `JWTClaims<T>` if `T` is Display.
pub trait DisplayJWTClaimsExt<T> {
fn display_jwt_claims(&self) -> DisplayJWTClaims<T>;
}

impl<T> DisplayJWTClaimsExt<T> for JWTClaims<T>
where T: fmt::Display
{
fn display_jwt_claims(&self) -> DisplayJWTClaims<T> {
DisplayJWTClaims { claims: self }
}
}

#[cfg(test)]
mod tests {
use jwt_simple::prelude::coarsetime;

use crate::display_jwt_claims::DisplayJWTClaimsExt;

#[test]
fn test_display_jwt_claims() {
use jwt_simple::claims::JWTClaims;

let claims = JWTClaims {
issuer: Some("issuer".to_string()),
subject: None,
audiences: None,
jwt_id: None,
issued_at: Some(coarsetime::UnixTimeStamp::from_millis(1723102819023)),
expires_at: Some(coarsetime::UnixTimeStamp::from_millis(1723102819023)),
custom: "custom".to_string(),
invalid_before: None,
nonce: None,
};

let display = claims.display_jwt_claims();
assert_eq!(
format!("{}", display),
"JWTClaims{issuer: issuer, issued_at: 2024-08-08T07:40:19.022460Z+0000, expires_at: 2024-08-08T07:40:19.022460Z+0000, custom: custom}"
);
}
}
1 change: 1 addition & 0 deletions src/common/license/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,5 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod display_jwt_claims;
pub mod license;
pub mod license_manager;
118 changes: 95 additions & 23 deletions src/common/license/src/license.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,10 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt;

use databend_common_base::display::display_option::DisplayOptionExt;
use databend_common_base::display::display_slice::DisplaySliceExt;
use serde::Deserialize;
use serde::Serialize;

Expand Down Expand Up @@ -75,8 +76,8 @@ pub enum Feature {
Unknown,
}

impl Display for Feature {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
impl fmt::Display for Feature {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Feature::LicenseInfo => write!(f, "license_info"),
Feature::Vacuum => write!(f, "vacuum"),
Expand All @@ -93,23 +94,28 @@ impl Display for Feature {
Feature::ComputeQuota(v) => {
write!(f, "compute_quota(")?;

write!(f, "threads_num: ")?;
match &v.threads_num {
None => write!(f, "threads_num: unlimited,")?,
Some(threads_num) => write!(f, "threads_num: {}", *threads_num)?,
None => write!(f, "unlimited,")?,
Some(threads_num) => write!(f, "{}", *threads_num)?,
};

write!(f, ", memory_usage: ")?;
match v.memory_usage {
None => write!(f, "memory_usage: unlimited,"),
Some(memory_usage) => write!(f, "memory_usage: {}", memory_usage),
None => write!(f, "unlimited,")?,
Some(memory_usage) => write!(f, "{}", memory_usage)?,
}
write!(f, ")")
}
Feature::StorageQuota(v) => {
write!(f, "storage_quota(")?;

write!(f, "storage_usage: ")?;
match v.storage_usage {
None => write!(f, "storage_usage: unlimited,"),
Some(storage_usage) => write!(f, "storage_usage: {}", storage_usage),
None => write!(f, "unlimited,")?,
Some(storage_usage) => write!(f, "{}", storage_usage)?,
}
write!(f, ")")
}
Feature::AmendTable => write!(f, "amend_table"),
Feature::Unknown => write!(f, "unknown"),
Expand Down Expand Up @@ -170,23 +176,53 @@ pub struct LicenseInfo {
pub features: Option<Vec<Feature>>,
}

impl fmt::Display for LicenseInfo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"LicenseInfo{{ type: {}, org: {}, tenants: {}, features: [{}] }}",
self.r#type.display(),
self.org.display(),
self.tenants
.as_ref()
.map(|x| x.as_slice().display())
.display(),
self.display_features()
)
}
}

impl LicenseInfo {
pub fn display_features(&self) -> String {
// sort all features in alphabet order and ignore test feature
let mut features = self.features.clone().unwrap_or_default();
pub fn display_features(&self) -> impl fmt::Display + '_ {
/// sort all features in alphabet order and ignore test feature
struct DisplayFeatures<'a>(&'a LicenseInfo);

if features.is_empty() {
return String::from("Unlimited");
}
impl<'a> fmt::Display for DisplayFeatures<'a> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let Some(features) = self.0.features.clone() else {
return write!(f, "Unlimited");
};

features.sort();
let mut features = features
.into_iter()
.filter(|f| f != &Feature::Test)
.map(|x| x.to_string())
.collect::<Vec<_>>();

features
.iter()
.filter(|f| **f != Feature::Test)
.map(|f| f.to_string())
.collect::<Vec<_>>()
.join(",")
features.sort();

for (i, feat) in features.into_iter().enumerate() {
if i > 0 {
write!(f, ",")?;
}

write!(f, "{}", feat)?;
}
Ok(())
}
}

DisplayFeatures(self)
}

/// Get Storage Quota from given license info.
Expand Down Expand Up @@ -298,4 +334,40 @@ mod tests {
serde_json::from_str::<Feature>("\"ssss\"").unwrap()
);
}

#[test]
fn test_display_license_info() {
let license_info = LicenseInfo {
r#type: Some("enterprise".to_string()),
org: Some("databend".to_string()),
tenants: Some(vec!["databend_tenant".to_string(), "foo".to_string()]),
features: Some(vec![
Feature::LicenseInfo,
Feature::Vacuum,
Feature::Test,
Feature::VirtualColumn,
Feature::BackgroundService,
Feature::DataMask,
Feature::AggregateIndex,
Feature::InvertedIndex,
Feature::ComputedColumn,
Feature::StorageEncryption,
Feature::Stream,
Feature::AttacheTable,
Feature::ComputeQuota(ComputeQuota {
threads_num: Some(1),
memory_usage: Some(1),
}),
Feature::StorageQuota(StorageQuota {
storage_usage: Some(1),
}),
Feature::AmendTable,
]),
};

assert_eq!(
"LicenseInfo{ type: enterprise, org: databend, tenants: [databend_tenant,foo], features: [aggregate_index,amend_table,attach_table,background_service,compute_quota(threads_num: 1, memory_usage: 1),computed_column,data_mask,inverted_index,license_info,storage_encryption,storage_quota(storage_usage: 1),stream,vacuum,virtual_column] }",
license_info.to_string()
);
}
}
5 changes: 4 additions & 1 deletion src/query/service/src/table_functions/others/license_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,10 @@ impl LicenseInfoSource {
DataType::String,
Value::Scalar(Scalar::String(human_readable_available_time)),
),
BlockEntry::new(DataType::String, Value::Scalar(Scalar::String(feature_str))),
BlockEntry::new(
DataType::String,
Value::Scalar(Scalar::String(feature_str.to_string())),
),
],
1,
))
Expand Down

0 comments on commit 66c8dea

Please sign in to comment.