feat: http handler add endpoint /v1/login/. (#15483)

* refactor: improve code readability for DATABEND_COMMIT_VERSION.

* feat: http handler add endpoint `/v1/login/`.

Author: youngsofun

src/query/config/src/version.rs

@@ -17,17 +17,19 @@ use std::sync::LazyLock;
 use semver::Version;
 
 pub static DATABEND_COMMIT_VERSION: LazyLock<String> = LazyLock::new(|| {
-    let git_tag = option_env!("DATABEND_GIT_SEMVER");
+    let semver = option_env!("DATABEND_GIT_SEMVER");
     let git_sha = option_env!("VERGEN_GIT_SHA");
     let rustc_semver = option_env!("VERGEN_RUSTC_SEMVER");
     let timestamp = option_env!("VERGEN_BUILD_TIMESTAMP");
 
-    match (git_tag, git_sha, rustc_semver, timestamp) {
+    match (semver, git_sha, rustc_semver, timestamp) {
         #[cfg(not(feature = "simd"))]
-        (Some(v1), Some(v2), Some(v3), Some(v4)) => format!("{}-{}(rust-{}-{})", v1, v2, v3, v4),
+        (Some(semver), Some(git_sha), Some(rustc_semver), Some(timestamp)) => {
+            format!("{semver}-{git_sha}(rust-{rustc_semver}-{timestamp})")
+        }
         #[cfg(feature = "simd")]
-        (Some(v1), Some(v2), Some(v3), Some(v4)) => {
-            format!("{}-{}-simd(rust-{}-{})", v1, v2, v3, v4)
+        (Some(semver), Some(git_sha), Some(rustc_semver), Some(timestamp)) => {
+            format!("{semver}-{git_sha}-simd(rust-{rustc_semver}-{timestamp})")
         }
         _ => String::new(),
     }

src/query/service/src/servers/http/http_services.rs

@@ -29,6 +29,7 @@ use poem::listener::RustlsConfig;
 use poem::middleware::CatchPanic;
 use poem::middleware::NormalizePath;
 use poem::middleware::TrailingSlash;
+use poem::post;
 use poem::put;
 use poem::Endpoint;
 use poem::EndpointExt;
@@ -40,6 +41,7 @@ use crate::servers::http::middleware::HTTPSessionMiddleware;
 use crate::servers::http::middleware::PanicHandler;
 use crate::servers::http::v1::clickhouse_router;
 use crate::servers::http::v1::list_suggestions;
+use crate::servers::http::v1::login_handler;
 use crate::servers::http::v1::query_route;
 use crate::servers::http::v1::streaming_load;
 use crate::servers::Server;
@@ -96,6 +98,7 @@ impl HttpHandler {
     async fn build_router(&self, sock: SocketAddr) -> impl Endpoint {
         let ep_v1 = Route::new()
             .nest("/query", query_route())
+            .at("/login", post(login_handler))
             .at("/streaming_load", put(streaming_load))
             .at("/upload_to_stage", put(upload_to_stage))
             .at("/suggested_background_tasks", get(list_suggestions));

src/query/service/src/servers/http/middleware.rs

@@ -294,7 +294,7 @@ impl<E: Endpoint> Endpoint for HTTPSessionEndpoint<E> {
                     self.ep.call(req).await
                 }
                 Err(err) => match err.code() {
-                    ErrorCode::AUTHENTICATE_FAILURE => {
+                    ErrorCode::AUTHENTICATE_FAILURE | ErrorCode::UNKNOWN_USER => {
                         warn!(
                             "http auth failure: {method} {uri}, headers={:?}, error={}",
                             sanitize_request_headers(&headers),

src/query/service/src/servers/http/v1/http_query_handlers.rs

@@ -77,7 +77,7 @@ pub struct QueryError {
 }
 
 impl QueryError {
-    pub(crate) fn from_error_code(e: &ErrorCode) -> Self {
+    pub(crate) fn from_error_code(e: ErrorCode) -> Self {
         QueryError {
             code: e.code(),
             message: e.display_text(),
@@ -201,7 +201,7 @@ impl QueryResponse {
             stats_uri: Some(make_state_uri(&id)),
             final_uri: Some(make_final_uri(&id)),
             kill_uri: Some(make_kill_uri(&id)),
-            error: r.state.error.as_ref().map(QueryError::from_error_code),
+            error: r.state.error.map(QueryError::from_error_code),
             has_result_set: r.state.has_result_set,
         })
         .with_header(HEADER_QUERY_ID, id.clone())
@@ -382,7 +382,7 @@ pub(crate) async fn query_handler(
             Err(e) => {
                 error!("http query fail to start sql, error: {:?}", e);
                 ctx.set_fail();
-                Ok(req.fail_to_start_sql(&e).into_response())
+                Ok(req.fail_to_start_sql(e).into_response())
             }
         }
     }

src/query/service/src/servers/http/v1/login.rs

@@ -0,0 +1,91 @@
+// 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::collections::BTreeMap;
+
+use databend_common_config::QUERY_SEMVER;
+use databend_common_storages_fuse::TableContext;
+use jwt_simple::prelude::Deserialize;
+use jwt_simple::prelude::Serialize;
+use poem::error::Result as PoemResult;
+use poem::web::Json;
+use poem::IntoResponse;
+
+use crate::servers::http::v1::HttpQueryContext;
+use crate::servers::http::v1::QueryError;
+
+#[derive(Deserialize, Clone)]
+struct LoginRequest {
+    pub database: Option<String>,
+    pub role: Option<String>,
+    pub secondary_roles: Option<Vec<String>>,
+    pub settings: Option<BTreeMap<String, String>>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+struct LoginResponse {
+    pub version: String,
+    pub error: Option<QueryError>,
+}
+
+/// Although theses can be checked for each /v1/query for now,
+/// getting these error in `conn()` instead of `execute()` is more in line with what users expect.
+async fn check_login(
+    ctx: &HttpQueryContext,
+    req: &LoginRequest,
+) -> databend_common_exception::Result<()> {
+    let session = &ctx.session;
+    let table_ctx = session.create_query_context().await?;
+    if let Some(database) = &req.database {
+        let cat = session.get_current_catalog();
+        let cat = table_ctx.get_catalog(&cat).await?;
+        cat.get_database(&ctx.session.get_current_tenant(), database)
+            .await?;
+    }
+    if let Some(role_name) = &req.role {
+        session.set_current_role_checked(role_name).await?;
+    }
+    session
+        .set_secondary_roles_checked(req.secondary_roles.clone())
+        .await?;
+
+    if let Some(conf_settings) = &req.settings {
+        let settings = session.get_settings();
+        for (k, v) in conf_settings {
+            settings.set_setting(k.to_string(), v.to_string())?;
+        }
+    }
+    Ok(())
+}
+
+///  # For SQL driver implementer:
+/// - It is encouraged to call `/v1/login` when establishing connection, not mandatory for now.
+/// - May get 404 when talk to old server, may check `/health` (no `/v1` prefix) to ensure the host:port is not wrong.
+///
+///  # TODO (need design):
+/// - (optional) check client version.
+/// - Return token for auth in the following queries from this session, to make it a real login.
+#[poem::handler]
+#[async_backtrace::framed]
+pub(crate) async fn login_handler(
+    ctx: &HttpQueryContext,
+    Json(req): Json<LoginRequest>,
+) -> PoemResult<impl IntoResponse> {
+    let version = QUERY_SEMVER.to_string();
+    let error = check_login(ctx, &req)
+        .await
+        .map_err(QueryError::from_error_code)
+        .err();
+    Ok(Json(LoginResponse { version, error }))
+}

src/query/service/src/servers/http/v1/mod.rs

@@ -15,6 +15,7 @@
 mod http_query_handlers;
 pub mod json_block;
 mod load;
+mod login;
 mod query;
 mod stage;
 mod suggestions;
@@ -29,6 +30,7 @@ pub use http_query_handlers::QueryStats;
 pub(crate) use json_block::JsonBlock;
 pub use load::streaming_load;
 pub use load::LoadResponse;
+pub(crate) use login::login_handler;
 pub use query::ExecuteStateKind;
 pub use query::ExpiringMap;
 pub use query::ExpiringState;

src/query/service/src/servers/http/v1/query/http_query.rs

@@ -83,7 +83,7 @@ pub struct HttpQueryRequest {
 }
 
 impl HttpQueryRequest {
-    pub(crate) fn fail_to_start_sql(&self, err: &ErrorCode) -> impl IntoResponse {
+    pub(crate) fn fail_to_start_sql(&self, err: ErrorCode) -> impl IntoResponse {
         metrics_incr_http_response_errors_count(err.name(), err.code());
         let session = self.session.as_ref().map(|s| {
             let txn_state = if matches!(s.txn_state, Some(TxnState::Active)) {

src/query/service/src/servers/http/v1/query/http_query_context.rs

@@ -25,7 +25,7 @@ use crate::sessions::SessionManager;
 use crate::sessions::SessionType;
 
 pub struct HttpQueryContext {
-    session: Arc<Session>,
+    pub session: Arc<Session>,
     pub query_id: String,
     pub node_id: String,
     pub deduplicate_label: Option<String>,

tests/suites/1_stateful/09_http_handler/09_0004_login.result

@@ -0,0 +1,14 @@
+>>>> drop table if exists t1;
+>>>> drop table if exists t2;
+>>>> create table t1 (a int);
+# auth fail
+{"code":"401","message":"User 'user1'@'%' does not exist."}
+# empty body
+{"code":"400","message":"parse error: EOF while parsing a value at line 1 column 0"}
+# db
+{"code":1003,"message":"Unknown database 't1'","detail":""}
+# db not exists
+{"code":1003,"message":"Unknown database 't2'","detail":""}
+# allow unknown key
+null
+>>>> drop table if exists t1;

tests/suites/1_stateful/09_http_handler/09_0004_login.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+stmt "drop table if exists t1;"
+stmt "drop table if exists t2;"
+stmt "create table t1 (a int);"
+
+echo "# auth fail"
+curl -s --header 'Content-Type: application/json'  --request POST '127.0.0.1:8000/v1/login/'  --data-raw '{}' -u user1: | jq -c ".error"
+echo "# empty body"
+curl -s --header 'Content-Type: application/json'  --request POST '127.0.0.1:8000/v1/login/'  -u root: | jq -c ".error"
+echo "# db"
+curl -s --header 'Content-Type: application/json'  --request POST '127.0.0.1:8000/v1/login/'  --data-raw '{"database":"t1"}' -u root: | jq -c ".error"
+echo "# db not exists"
+curl -s --header 'Content-Type: application/json'  --request POST '127.0.0.1:8000/v1/login/'  --data-raw '{"database":"t2"}' -u root: | jq -c ".error"
+echo "# allow unknown key"
+curl -s --header 'Content-Type: application/json'  --request POST '127.0.0.1:8000/v1/login/'  --data-raw '{"unknown": null}' -u root:xx | jq -c ".error"
+
+stmt "drop table if exists t1;"