feat: auth by refresh and session tokens.

src/common/exception/src/exception_code.rs

@@ -413,4 +413,9 @@ build_exceptions! {
 build_exceptions! {
     // A task that already stopped and can not stopped twice.
     AlreadyStopped(5002),
+
+    SessionTokenExpired(5100),
+    RefreshTokenExpired(5101),
+    SessionTokenNotFound(5102),
+    RefreshTokenNotFound(5103)
 }

src/meta/app/src/principal/mod.rs

@@ -46,6 +46,8 @@ pub mod tenant_user_ident;
 pub mod user_defined_file_format_ident;
 pub mod user_setting_ident;
 pub mod user_stage_ident;
+pub mod user_token;
+pub mod user_token_ident;
 
 pub use connection::*;
 pub use file_format::*;

src/meta/app/src/principal/user_token.rs

@@ -0,0 +1,35 @@
+// 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 serde::Deserialize;
+use serde::Serialize;
+
+/// A client starts with /session/login to get the initial refresh_token and session_token pair.
+/// - Use session_token for computing.
+/// - Use refresh_token to auth /session/renew and get new pair when session_token expires.
+#[derive(
+    serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, num_derive::FromPrimitive,
+)]
+pub enum TokenType {
+    Refresh = 1,
+    Session = 2,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
+pub struct QueryTokenInfo {
+    pub token_type: TokenType,
+    /// used to delete refresh token when close session, which authed by session_token too.
+    /// None for Refresh token.
+    pub parent: Option<String>,
+}

src/meta/app/src/principal/user_token_ident.rs

@@ -0,0 +1,60 @@
+// 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 crate::tenant_key::ident::TIdent;
+
+/// Define the meta-service key for a user setting.
+pub type TokenIdent = TIdent<Resource>;
+
+pub use kvapi_impl::Resource;
+
+mod kvapi_impl {
+
+    use databend_common_meta_kvapi::kvapi;
+
+    use crate::principal::user_token::QueryTokenInfo;
+    use crate::tenant_key::resource::TenantResource;
+
+    pub struct Resource;
+    impl TenantResource for Resource {
+        const PREFIX: &'static str = "__fd_token";
+        const TYPE: &'static str = "TokenIdent";
+        const HAS_TENANT: bool = true;
+        type ValueType = QueryTokenInfo;
+    }
+
+    impl kvapi::Value for QueryTokenInfo {
+        fn dependency_keys(&self) -> impl IntoIterator<Item = String> {
+            []
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use databend_common_meta_kvapi::kvapi::Key;
+
+    use crate::principal::user_token_ident::TokenIdent;
+    use crate::tenant::Tenant;
+
+    #[test]
+    fn test_setting_ident() {
+        let tenant = Tenant::new_literal("tenant1");
+        let ident = TokenIdent::new(tenant.clone(), "test");
+        assert_eq!("__fd_token/tenant1/test", ident.to_string_key());
+
+        let got = TokenIdent::from_str_key(&ident.to_string_key()).unwrap();
+        assert_eq!(ident, got);
+    }
+}

src/meta/proto-conv/src/lib.rs

@@ -85,6 +85,7 @@ mod stage_from_to_protobuf_impl;
 mod table_from_to_protobuf_impl;
 mod tenant_quota_from_to_protobuf_impl;
 mod tident_from_to_protobuf_impl;
+mod token_from_to_protobuf_impl;
 mod udf_from_to_protobuf_impl;
 mod user_from_to_protobuf_impl;
 mod util;

src/meta/proto-conv/src/token_from_to_protobuf_impl.rs

@@ -0,0 +1,57 @@
+// 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.
+
+//! This mod is the key point about compatibility.
+//! Everytime update anything in this file, update the `VER` and let the tests pass.
+
+use databend_common_meta_app::principal::user_token as mt;
+use databend_common_protos::pb;
+use num::FromPrimitive;
+
+use crate::reader_check_msg;
+use crate::FromToProto;
+use crate::Incompatible;
+use crate::MIN_READER_VER;
+use crate::VER;
+
+impl FromToProto for mt::QueryTokenInfo {
+    type PB = pb::TokenInfo;
+
+    fn get_pb_ver(p: &Self::PB) -> u64 {
+        p.ver
+    }
+
+    fn from_pb(p: Self::PB) -> Result<Self, Incompatible>
+    where Self: Sized {
+        reader_check_msg(p.ver, p.min_reader_ver)?;
+
+        let v = Self {
+            token_type: FromPrimitive::from_i32(p.token_type).ok_or_else(|| Incompatible {
+                reason: format!("invalid TokenType: {}", p.token_type),
+            })?,
+            parent: p.parent,
+        };
+        Ok(v)
+    }
+
+    fn to_pb(&self) -> Result<Self::PB, Incompatible> {
+        let p = pb::TokenInfo {
+            ver: VER,
+            min_reader_ver: MIN_READER_VER,
+            token_type: self.token_type.clone() as i32,
+            parent: self.parent.clone(),
+        };
+        Ok(p)
+    }
+}

src/meta/proto-conv/src/util.rs

@@ -134,6 +134,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (102, "2024-07-11: Add: UserOption add must_change_password, AuthInfo.Password add need_change"),
     (103, "2024-07-31: Add: ShareMetaV2"),
     (104, "2024-08-02: Add: add share catalog into Catalog meta"),
+    (105, "2024-08-08: Add: add QueryTokenInfo"),
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)

src/meta/proto-conv/tests/it/main.rs

@@ -108,3 +108,4 @@ mod v101_database_meta;
 mod v102_user_must_change_password;
 mod v103_share_meta_v2;
 mod v104_share_catalog;
+mod v105_query_token;

src/meta/proto-conv/tests/it/v105_query_token.rs

@@ -0,0 +1,44 @@
+// Copyright 2023 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 databend_common_meta_app::principal::user_token::TokenType;
+use fastrace::func_name;
+
+use crate::common;
+
+// These bytes are built when a new version in introduced,
+// and are kept for backward compatibility test.
+//
+// *************************************************************
+// * These messages should never be updated,                   *
+// * only be added when a new version is added,                *
+// * or be removed when an old version is no longer supported. *
+// *************************************************************
+//
+// The message bytes are built from the output of `test_pb_from_to()`
+#[test]
+fn test_v105_query_token_info() -> anyhow::Result<()> {
+    let query_token_info_v105 = vec![
+        8, 1, 18, 17, 112, 97, 114, 101, 110, 116, 95, 116, 111, 107, 101, 110, 95, 104, 97, 115,
+        104, 160, 6, 105, 168, 6, 24,
+    ];
+
+    let want = || databend_common_meta_app::principal::user_token::QueryTokenInfo {
+        token_type: TokenType::Refresh,
+        parent: Some("parent_token_hash".to_string()),
+    };
+
+    common::test_pb_from_to(func_name!(), want())?;
+    common::test_load_old(func_name!(), query_token_info_v105.as_slice(), 105, want())
+}

src/meta/protos/proto/token.proto

@@ -0,0 +1,33 @@
+// 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.
+
+// The identifier of a database by name. Names can be changed.
+// There is no guarantee that two get-database request by name will return the
+// same instance.
+
+syntax = "proto3";
+
+package databend_proto;
+
+message TokenInfo {
+  uint64 ver = 100;
+  uint64 min_reader_ver = 101;
+
+  enum TokenType {
+    Session = 0;
+    Refresh = 2;
+  }
+  TokenType token_type = 1;
+  optional string parent = 2;
+}

src/query/management/src/lib.rs

@@ -28,6 +28,7 @@ pub mod udf;
 mod user;
 
 pub mod errors;
+mod token;
 
 pub use cluster::ClusterApi;
 pub use cluster::ClusterMgr;
@@ -46,5 +47,6 @@ pub use setting::SettingApi;
 pub use setting::SettingMgr;
 pub use stage::StageApi;
 pub use stage::StageMgr;
+pub use token::TokenMgr;
 pub use user::UserApi;
 pub use user::UserMgr;

src/query/management/src/token.rs

@@ -0,0 +1,97 @@
+// 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::sync::Arc;
+use std::time::Duration;
+
+use databend_common_exception::Result;
+use databend_common_meta_api::kv_pb_api::KVPbApi;
+use databend_common_meta_api::kv_pb_api::UpsertPB;
+use databend_common_meta_app::principal::user_token::QueryTokenInfo;
+use databend_common_meta_app::principal::user_token_ident::TokenIdent;
+use databend_common_meta_app::tenant::Tenant;
+use databend_common_meta_kvapi::kvapi;
+use databend_common_meta_kvapi::kvapi::Key;
+use databend_common_meta_kvapi::kvapi::UpsertKVReq;
+use databend_common_meta_types::MatchSeq;
+use databend_common_meta_types::MetaError;
+use databend_common_meta_types::Operation;
+use databend_common_meta_types::With;
+
+pub struct TokenMgr {
+    kv_api: Arc<dyn kvapi::KVApi<Error = MetaError>>,
+    tenant: Tenant,
+}
+
+impl TokenMgr {
+    pub fn create(kv_api: Arc<dyn kvapi::KVApi<Error = MetaError>>, tenant: &Tenant) -> Self {
+        TokenMgr {
+            kv_api,
+            tenant: tenant.clone(),
+        }
+    }
+
+    fn token_ident(&self, token_hash: &str) -> TokenIdent {
+        TokenIdent::new(self.tenant.clone(), token_hash)
+    }
+}
+
+impl TokenMgr {
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    pub async fn upsert_token(
+        &self,
+        token_hash: &str,
+        token_info: QueryTokenInfo,
+        ttl: Duration,
+        update_only: bool,
+    ) -> Result<bool> {
+        let ident = self.token_ident(token_hash);
+        let seq = MatchSeq::GE(if update_only { 1 } else { 0 });
+        let upsert = UpsertPB::update(ident, token_info)
+            .with(seq)
+            .with_ttl(Duration::from_secs(ttl.as_secs()));
+
+        let res = self.kv_api.upsert_pb(&upsert).await?;
+
+        Ok(res.prev.is_none())
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    pub async fn get_token(&self, token_hash: &str) -> Result<Option<QueryTokenInfo>> {
+        let ident = self.token_ident(token_hash);
+        let res = self.kv_api.get_pb(&ident).await?;
+
+        Ok(res.map(|r| r.data))
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    pub async fn drop_token(&self, token_hash: &str) -> Result<()> {
+        let key = self.token_ident(token_hash).to_string_key();
+
+        // simply ignore the result
+        self.kv_api
+            .upsert_kv(UpsertKVReq::new(
+                &key,
+                MatchSeq::GE(0),
+                Operation::Delete,
+                None,
+            ))
+            .await?;
+
+        Ok(())
+    }
+}