Skip to content

Commit

Permalink
fix: exit client if login failed with password wrong (#474)
Browse files Browse the repository at this point in the history
* fix: exit client if login failed with password wrong

* bump jsonb 0.4.1
  • Loading branch information
b41sh authored Aug 15, 2024
1 parent b590fe1 commit c2f7e60
Show file tree
Hide file tree
Showing 8 changed files with 75 additions and 18 deletions.
6 changes: 4 additions & 2 deletions bindings/nodejs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ impl RowIterator {
/// Fetch next row.
/// Returns `None` if there are no more rows.
#[napi]
#[allow(clippy::missing_safety_doc)]
pub async unsafe fn next(&mut self) -> Option<Result<Row>> {
self.0
.next()
Expand All @@ -206,6 +207,7 @@ impl RowIteratorExt {
/// Fetch next row or stats.
/// Returns `None` if there are no more rows.
#[napi]
#[allow(clippy::missing_safety_doc)]
pub async unsafe fn next(&mut self) -> Option<Result<RowOrStats>> {
match self.0.next().await {
None => None,
Expand Down Expand Up @@ -322,7 +324,7 @@ impl Client {
self.0
.get_conn()
.await
.map(|conn| Connection(conn))
.map(Connection)
.map_err(format_napi_error)
}
}
Expand Down Expand Up @@ -366,7 +368,7 @@ impl Connection {
.await
.map_err(format_napi_error)?
.into_iter()
.map(|row| Row(row))
.map(Row)
.collect())
}

Expand Down
5 changes: 1 addition & 4 deletions bindings/python/src/blocking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,7 @@ pub struct BlockingDatabendConnection(Box<dyn databend_driver::Connection>);
impl BlockingDatabendConnection {
pub fn info(&self, py: Python) -> PyResult<ConnectionInfo> {
let this = self.0.clone();
let ret = wait_for_future(py, async move {
let info = this.info().await;
info
});
let ret = wait_for_future(py, async move { this.info().await });
Ok(ConnectionInfo::new(ret))
}

Expand Down
11 changes: 2 additions & 9 deletions bindings/python/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -162,10 +162,7 @@ impl RowIterator {
impl RowIterator {
pub fn schema(&self, py: Python) -> PyResult<Schema> {
let streamer = self.0.clone();
let ret = wait_for_future(py, async move {
let schema = streamer.lock().await.schema();
schema
});
let ret = wait_for_future(py, async move { streamer.lock().await.schema() });
Ok(Schema(ret))
}

Expand Down Expand Up @@ -208,11 +205,7 @@ pub struct Schema(databend_driver::SchemaRef);
#[pymethods]
impl Schema {
pub fn fields<'p>(&'p self, py: Python<'p>) -> PyResult<Bound<'p, PyList>> {
let fields = self
.0
.fields()
.into_iter()
.map(|f| Field(f.clone()).into_py(py));
let fields = self.0.fields().iter().map(|f| Field(f.clone()).into_py(py));
Ok(PyList::new_bound(py, fields))
}
}
Expand Down
1 change: 1 addition & 0 deletions cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ databend-driver = { workspace = true, features = ["rustls", "flight-sql"] }
tokio-stream = { workspace = true }

anyhow = "1.0"
arrow = { workspace = true }
async-recursion = "1.1.0"
async-trait = "0.1"
clap = { version = "4.4", features = ["derive", "env"] }
Expand Down
28 changes: 27 additions & 1 deletion cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -330,6 +330,7 @@ pub async fn main() -> Result<()> {
}
}

let user = conn_args.user.clone();
let dsn = conn_args.get_dsn()?;
let mut settings = Settings::default();
let is_terminal = stdin().is_terminal();
Expand Down Expand Up @@ -365,7 +366,32 @@ pub async fn main() -> Result<()> {
}
settings.time = args.time;

let mut session = session::Session::try_new(dsn, settings, is_repl).await?;
let mut session = match session::Session::try_new(dsn, settings, is_repl).await {
Ok(session) => session,
Err(err) => {
// Exit client if user login failed.
if let Some(error) = err.downcast_ref::<databend_driver::Error>() {
match error {
databend_driver::Error::Api(
databend_client::error::Error::InvalidResponse(resp_err),
) => {
if resp_err.code == 401 {
println!("Authenticate failed wrong password user {}", user);
return Ok(());
}
}
databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(ipc_err)) => {
if ipc_err.contains("Unauthenticated") {
println!("Authenticate failed wrong password user {}", user);
return Ok(());
}
}
_ => {}
}
}
return Err(err);
}
};

let log_dir = format!(
"{}/.bendsql",
Expand Down
24 changes: 23 additions & 1 deletion cli/src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,29 @@ impl Session {
);
}
}
let version = conn.version().await.unwrap_or_default();
let version = match conn.version().await {
Ok(version) => version,
Err(err) => {
match err {
databend_driver::Error::Api(
databend_client::error::Error::InvalidResponse(ref resp_err),
) => {
if resp_err.code == 401 {
return Err(err.into());
}
}
databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(
ref ipc_err,
)) => {
if ipc_err.contains("Unauthenticated") {
return Err(err.into());
}
}
_ => {}
}
"".to_string()
}
};
println!("Connected to {}", version);

let config = sled::Config::new().temporary(true);
Expand Down
16 changes: 16 additions & 0 deletions core/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,14 @@ impl APIClient {
resp = builder.headers(headers.clone()).send().await?;
}
if resp.status() != 200 {
if resp.status() == 401 {
let resp_err = QueryError {
code: resp.status().as_u16(),
message: resp.text().await.unwrap_or_default(),
detail: None,
};
return Err(Error::InvalidResponse(resp_err));
}
return Err(Error::Request(format!(
"Start Query failed with status {}: {}",
resp.status(),
Expand Down Expand Up @@ -357,6 +365,14 @@ impl APIClient {
if resp.status() == 404 {
return Err(Error::SessionTimeout(resp.text().await?));
}
if resp.status() == 401 {
let resp_err = QueryError {
code: resp.status().as_u16(),
message: resp.text().await.unwrap_or_default(),
detail: None,
};
return Err(Error::InvalidResponse(resp_err));
}
return Err(Error::Request(format!(
"Query Page failed with status {}: {}",
resp.status(),
Expand Down
2 changes: 1 addition & 1 deletion sql/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ geozero = { version = "0.13", features = ["default", "with-wkb"] }
glob = "0.3"
hex = "0.4.3"
itertools = "0.12"
jsonb = "0.4"
jsonb = "0.4.1"
lexical-core = "0.8"
memchr = "2.7"
roaring = { version = "0.10", features = ["serde"] }
Expand Down

0 comments on commit c2f7e60

Please sign in to comment.