chore: impl Display for LicenseInfo and JWTClaims<T> (#16244)

Author: drmingdrmer

src/common/base/src/display/display_slice.rs

@@ -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());
+    }
+}

src/common/base/src/display/mod.rs

@@ -13,4 +13,5 @@
 // limitations under the License.
 
 pub mod display_option;
+pub mod display_slice;
 pub mod display_unix_epoch;

src/common/license/src/display_jwt_claims.rs

@@ -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}"
+        );
+    }
+}

src/common/license/src/lib.rs

@@ -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;

src/common/license/src/license.rs

@@ -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;
 
@@ -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"),
@@ -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"),
@@ -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.
@@ -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()
+        );
+    }
 }

src/query/service/src/table_functions/others/license_info.rs

@@ -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,
         ))