dev

Author: MikeNitsenko

packages/cubejs-backend-native/src/auth.rs

@@ -1,7 +1,7 @@
 use async_trait::async_trait;
 use cubesql::{
     di_service,
-    sql::{AuthContext, AuthenticateResponse, SqlAuthService},
+    sql::{AuthContext, AuthenticateResponse, SqlAuthService, SqlAuthServiceAuthenticateRequest},
     transport::LoadRequestMeta,
     CubeError,
 };
@@ -50,9 +50,28 @@ pub struct TransportRequest {
     pub meta: Option<LoadRequestMeta>,
 }
 
+#[derive(Debug, Serialize)]
+pub struct TransportAuthRequest {
+    pub id: String,
+    pub meta: Option<LoadRequestMeta>,
+    pub protocol: String,
+    pub method: String,
+}
+
+impl From<(TransportRequest, SqlAuthServiceAuthenticateRequest)> for TransportAuthRequest {
+    fn from((t, a): (TransportRequest, SqlAuthServiceAuthenticateRequest)) -> Self {
+        Self {
+            id: t.id,
+            meta: t.meta,
+            protocol: a.protocol,
+            method: a.method,
+        }
+    }
+}
+
 #[derive(Debug, Serialize)]
 struct CheckSQLAuthTransportRequest {
-    request: TransportRequest,
+    request: TransportAuthRequest,
     user: Option<String>,
     password: Option<String>,
 }
@@ -92,6 +111,7 @@ impl AuthContext for NativeSQLAuthContext {
 impl SqlAuthService for NodeBridgeAuthService {
     async fn authenticate(
         &self,
+        sql_auth_request: SqlAuthServiceAuthenticateRequest,
         user: Option<String>,
         password: Option<String>,
     ) -> Result<AuthenticateResponse, CubeError> {
@@ -100,9 +120,11 @@ impl SqlAuthService for NodeBridgeAuthService {
         let request_id = Uuid::new_v4().to_string();
 
         let extra = serde_json::to_string(&CheckSQLAuthTransportRequest {
-            request: TransportRequest {
+            request: TransportAuthRequest {
                 id: format!("{}-span-1", request_id),
                 meta: None,
+                protocol: sql_auth_request.protocol.clone(),
+                method: sql_auth_request.method.clone(),
             },
             user: user.clone(),
             password: password.clone(),

rust/cubesql/cubesql/src/compile/router.rs

@@ -12,6 +12,7 @@ use crate::{
         DatabaseVariable, DatabaseVariablesToUpdate,
     },
     sql::{
+        auth_service::SqlAuthServiceAuthenticateRequest,
         dataframe,
         statement::{
             ApproximateCountDistinctVisitor, CastReplacer, DateTokenNormalizeReplacer,
@@ -447,12 +448,16 @@ impl QueryRouter {
                 })?
             {
                 self.state.set_user(Some(to_user.clone()));
+                let sql_auth_request = SqlAuthServiceAuthenticateRequest {
+                    protocol: "postgres".to_string(),
+                    method: "password".to_string(),
+                };
                 let authenticate_response = self
                     .session_manager
                     .server
                     .auth
                     // TODO do we want to send actual password here?
-                    .authenticate(Some(to_user.clone()), None)
+                    .authenticate(sql_auth_request, Some(to_user.clone()), None)
                     .await
                     .map_err(|e| {
                         CompilationError::internal(format!("Error calling authenticate: {}", e))
@@ -562,11 +567,15 @@ impl QueryRouter {
 
     async fn reauthenticate_if_needed(&self) -> CompilationResult<()> {
         if self.state.is_auth_context_expired() {
+            let sql_auth_request = SqlAuthServiceAuthenticateRequest {
+                protocol: "postgres".to_string(),
+                method: "password".to_string(),
+            };
             let authenticate_response = self
                 .session_manager
                 .server
                 .auth
-                .authenticate(self.state.user(), None)
+                .authenticate(sql_auth_request, self.state.user(), None)
                 .await
                 .map_err(|e| {
                     CompilationError::fatal(format!(

rust/cubesql/cubesql/src/compile/test/mod.rs

@@ -8,9 +8,10 @@ use crate::{
     },
     config::{ConfigObj, ConfigObjImpl},
     sql::{
-        compiler_cache::CompilerCacheImpl, dataframe::batches_to_dataframe,
-        pg_auth_service::PostgresAuthServiceDefaultImpl, AuthContextRef, AuthenticateResponse,
-        HttpAuthContext, ServerManager, Session, SessionManager, SqlAuthService,
+        auth_service::SqlAuthServiceAuthenticateRequest, compiler_cache::CompilerCacheImpl,
+        dataframe::batches_to_dataframe, pg_auth_service::PostgresAuthServiceDefaultImpl,
+        AuthContextRef, AuthenticateResponse, HttpAuthContext, ServerManager, Session,
+        SessionManager, SqlAuthService,
     },
     transport::{
         CubeMeta, CubeMetaDimension, CubeMetaJoin, CubeMetaMeasure, CubeMetaSegment,
@@ -750,6 +751,7 @@ pub fn get_test_auth() -> Arc<dyn SqlAuthService> {
     impl SqlAuthService for TestSqlAuth {
         async fn authenticate(
             &self,
+            _request: SqlAuthServiceAuthenticateRequest,
             _user: Option<String>,
             password: Option<String>,
         ) -> Result<AuthenticateResponse, CubeError> {

rust/cubesql/cubesql/src/sql/auth_service.rs

@@ -2,6 +2,7 @@ use std::{any::Any, env, fmt::Debug, sync::Arc};
 
 use crate::CubeError;
 use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
 // We cannot use generic here. It's why there is this trait
@@ -43,10 +44,17 @@ pub struct AuthenticateResponse {
     pub skip_password_check: bool,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SqlAuthServiceAuthenticateRequest {
+    pub protocol: String,
+    pub method: String,
+}
+
 #[async_trait]
 pub trait SqlAuthService: Send + Sync + Debug {
     async fn authenticate(
         &self,
+        sql_auth_request: SqlAuthServiceAuthenticateRequest,
         user: Option<String>,
         password: Option<String>,
     ) -> Result<AuthenticateResponse, CubeError>;
@@ -61,6 +69,7 @@ crate::di_service!(SqlAuthDefaultImpl, [SqlAuthService]);
 impl SqlAuthService for SqlAuthDefaultImpl {
     async fn authenticate(
         &self,
+        _request: SqlAuthServiceAuthenticateRequest,
         _user: Option<String>,
         password: Option<String>,
     ) -> Result<AuthenticateResponse, CubeError> {

rust/cubesql/cubesql/src/sql/mod.rs

@@ -13,7 +13,7 @@ pub(crate) mod types;
 // Public API
 pub use auth_service::{
     AuthContext, AuthContextRef, AuthenticateResponse, HttpAuthContext, SqlAuthDefaultImpl,
-    SqlAuthService,
+    SqlAuthService, SqlAuthServiceAuthenticateRequest,
 };
 pub use database_variables::postgres::session_vars::CUBESQL_PENALIZE_POST_PROCESSING_VAR;
 pub use postgres::*;

rust/cubesql/cubesql/src/sql/postgres/pg_auth_service.rs

@@ -3,6 +3,7 @@ use std::{collections::HashMap, fmt::Debug, sync::Arc};
 use async_trait::async_trait;
 
 use crate::{
+    sql::auth_service::SqlAuthServiceAuthenticateRequest,
     sql::{AuthContextRef, SqlAuthService},
     CubeError,
 };
@@ -74,8 +75,16 @@ impl PostgresAuthService for PostgresAuthServiceDefaultImpl {
         }
 
         let user = parameters.get("user").unwrap().clone();
+        let sql_auth_request = SqlAuthServiceAuthenticateRequest {
+            protocol: "postgres".to_string(),
+            method: "password".to_string(),
+        };
         let authenticate_response = service
-            .authenticate(Some(user.clone()), Some(password_message.password.clone()))
+            .authenticate(
+                sql_auth_request,
+                Some(user.clone()),
+                Some(password_message.password.clone()),
+            )
             .await;
 
         let auth_fail = || {