refactor: some http APIs do not need sticky. (#16644)

* refactor: some http APIs do not need sticky. * fix forward log.
databendlabs · Oct 20, 2024 · c88360c · c88360c
1 parent 4557131
commit c88360c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/src/query/service/src/servers/http/middleware/session.rs b/src/query/service/src/servers/http/middleware/session.rs
@@ -79,6 +79,8 @@ pub enum EndpointKind {
     Clickhouse,
     NoAuth,
     Verify,
+    UploadToStage,
+    SystemInfo,
 }
 
 impl EndpointKind {
@@ -94,18 +96,22 @@ impl EndpointKind {
     }
     pub fn require_databend_token_type(&self) -> Result<Option<TokenType>> {
         match self {
-            EndpointKind::Verify => Ok(None),
+            EndpointKind::Verify | EndpointKind::NoAuth => Ok(None),
             EndpointKind::Refresh => Ok(Some(TokenType::Refresh)),
-            EndpointKind::StartQuery | EndpointKind::PollQuery | EndpointKind::Logout => {
+            EndpointKind::StartQuery
+            | EndpointKind::PollQuery
+            | EndpointKind::Logout
+            | EndpointKind::SystemInfo
+            | EndpointKind::UploadToStage => {
                 if GlobalConfig::instance().query.management_mode {
                     Ok(None)
                 } else {
                     Ok(Some(TokenType::Session))
                 }
             }
-            _ => Err(ErrorCode::AuthenticateFailure(format!(
-                "should not use databend token for {self:?}",
-            ))),
+            EndpointKind::Login | EndpointKind::Clickhouse => Err(ErrorCode::AuthenticateFailure(
+                format!("should not use databend token for {self:?}",),
+            )),
         }
     }
 }
@@ -455,7 +461,7 @@ impl<E: Endpoint> Endpoint for HTTPSessionEndpoint<E> {
                     .map_err(HttpErrorCode::server_error)?
                 {
                     log::info!(
-                        "forwarding {} from {local_id} to {sticky_node_id}",
+                        "forwarding /v1{} from {local_id} to {sticky_node_id}",
                         req.uri()
                     );
                     forward_request(req, node).await

diff --git a/src/query/service/src/servers/http/v1/http_query_handlers.rs b/src/query/service/src/servers/http/v1/http_query_handlers.rs
@@ -450,17 +450,17 @@ pub fn query_route() -> Route {
         (
             "/upload_to_stage",
             put(upload_to_stage),
-            EndpointKind::StartQuery,
+            EndpointKind::UploadToStage,
         ),
         (
             "/suggested_background_tasks",
             get(list_suggestions),
-            EndpointKind::StartQuery,
+            EndpointKind::SystemInfo,
         ),
         (
             "/discovery_nodes",
             get(discovery_nodes),
-            EndpointKind::StartQuery,
+            EndpointKind::SystemInfo,
         ),
     ];