Skip to content

Commit 25c655e

Browse files
matthewmturnermatthewmturner
committed
More port args
1 parent 65e6cd1 commit 25c655e

6 files changed

Lines changed: 100 additions & 54 deletions

File tree

‎src/args.rs‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -132,24 +132,20 @@ pub enum Command {
132132
ServeHttp {
133133
#[clap(short, long)]
134134
config: Option<String>,
135-
#[clap(
136-
long,
137-
global = true,
138-
help = "Set the host and port to be used for server"
139-
)]
140-
host: Option<String>,
135+
#[clap(long, help = "Set the port to be used for server")]
136+
port: Option<String>,
137+
#[clap(long, help = "Set the port to be used for serving metrics")]
138+
metrics_port: Option<String>,
141139
},
142140
/// Start a FlightSQL server
143141
#[command(name = "serve-flightsql")]
144142
ServeFlightSql {
145143
#[clap(short, long)]
146144
config: Option<String>,
147-
#[clap(
148-
long,
149-
global = true,
150-
help = "Set the host and port to be used for server"
151-
)]
152-
host: Option<String>,
145+
#[clap(long, help = "Set the port to be used for server")]
146+
port: Option<String>,
147+
#[clap(long, help = "Set the port to be used for serving metrics")]
148+
metrics_port: Option<String>,
153149
},
154150
}
155151

‎src/server/flightsql/mod.rs‎

Lines changed: 37 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -188,21 +188,47 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
188188
execution_ctx.execute_ddl().await;
189189
}
190190
let app_execution = AppExecution::new(execution_ctx);
191-
let url = if let Some(cmd) = cli.command.clone() {
191+
192+
let (addr, metrics_addr) = if let Some(cmd) = cli.command.clone() {
192193
match cmd {
193-
Command::ServeFlightSql { host, .. } => host,
194-
_ => None,
194+
Command::ServeFlightSql {
195+
port: Some(port),
196+
metrics_port: Some(metrics_port),
197+
..
198+
} => (
199+
format!("localhost:{port}"),
200+
format!("0.0.0.0:{metrics_port}"),
201+
),
202+
Command::ServeFlightSql {
203+
port: Some(port),
204+
metrics_port: None,
205+
..
206+
} => (
207+
format!("localhost:{port}"),
208+
config.flightsql_server.server_metrics_port.clone(),
209+
),
210+
Command::ServeFlightSql {
211+
port: None,
212+
metrics_port: Some(metrics_port),
213+
..
214+
} => (
215+
DEFAULT_SERVER_ADDRESS.to_string(),
216+
format!("0.0.0.0:{metrics_port}"),
217+
),
218+
219+
_ => (
220+
DEFAULT_SERVER_ADDRESS.to_string(),
221+
config.flightsql_server.server_metrics_port.clone(),
222+
),
195223
}
196224
} else {
197-
None
225+
(
226+
DEFAULT_SERVER_ADDRESS.to_string(),
227+
config.flightsql_server.server_metrics_port.clone(),
228+
)
198229
};
199-
let app = FlightSqlApp::try_new(
200-
app_execution,
201-
&config,
202-
&url.unwrap_or(DEFAULT_SERVER_ADDRESS.to_string()),
203-
&config.flightsql_server.server_metrics_port,
204-
)
205-
.await?;
230+
231+
let app = FlightSqlApp::try_new(app_execution, &config, &addr, &metrics_addr).await?;
206232
app.run().await;
207233
Ok(())
208234
}

‎src/server/http/mod.rs‎

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -157,21 +157,45 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
157157
}
158158
}
159159
debug!("Created AppExecution: {app_execution:?}");
160-
let url = if let Some(cmd) = cli.command.clone() {
160+
let (addr, metrics_addr) = if let Some(cmd) = cli.command.clone() {
161161
match cmd {
162-
Command::ServeHttp { host, .. } => host,
163-
_ => None,
162+
Command::ServeHttp {
163+
port: Some(port),
164+
metrics_port: Some(metrics_port),
165+
..
166+
} => (
167+
format!("localhost:{port}"),
168+
format!("0.0.0.0:{metrics_port}"),
169+
),
170+
Command::ServeHttp {
171+
port: Some(port),
172+
metrics_port: None,
173+
..
174+
} => (
175+
format!("localhost:{port}"),
176+
config.http_server.server_metrics_port.clone(),
177+
),
178+
Command::ServeHttp {
179+
port: None,
180+
metrics_port: Some(metrics_port),
181+
..
182+
} => (
183+
DEFAULT_SERVER_ADDRESS.to_string(),
184+
format!("0.0.0.0:{metrics_port}"),
185+
),
186+
187+
_ => (
188+
DEFAULT_SERVER_ADDRESS.to_string(),
189+
config.http_server.server_metrics_port.clone(),
190+
),
164191
}
165192
} else {
166-
None
193+
(
194+
DEFAULT_SERVER_ADDRESS.to_string(),
195+
config.http_server.server_metrics_port.clone(),
196+
)
167197
};
168-
let app = HttpApp::try_new(
169-
app_execution,
170-
config.clone(),
171-
&url.unwrap_or(DEFAULT_SERVER_ADDRESS.to_string()),
172-
&config.http_server.server_metrics_port,
173-
)
174-
.await?;
198+
let app = HttpApp::try_new(app_execution, config.clone(), &addr, &metrics_addr).await?;
175199
app.run().await;
176200

177201
Ok(())

‎src/tui/handlers/mod.rs‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,6 @@ use log::{error, info, trace};
2424
use ratatui::crossterm::event::{self, KeyCode, KeyEvent};
2525
use tui_logger::TuiWidgetEvent;
2626

27-
#[cfg(feature = "flightsql")]
28-
use crate::args::Command;
2927
#[cfg(feature = "flightsql")]
3028
use crate::tui::state::tabs::flightsql::FlightSQLConnectionStatus;
3129
use crate::tui::state::tabs::history::Context;
@@ -247,16 +245,9 @@ pub fn app_event_handler(app: &mut App, event: AppEvent) -> Result<()> {
247245
AppEvent::FlightSQLEstablishConnection => {
248246
let execution = Arc::clone(&app.execution);
249247
let _event_tx = app.event_tx.clone();
250-
let cli_url = if let Some(cmd) = app.args.command.clone() {
251-
match cmd {
252-
Command::ServeFlightSql { host, .. } => host,
253-
Command::ServeHttp { host, .. } => host,
254-
}
255-
} else {
256-
None
257-
};
248+
let host = app.args.host.clone();
258249
tokio::spawn(async move {
259-
if let Err(e) = execution.create_flightsql_client(cli_url).await {
250+
if let Err(e) = execution.create_flightsql_client(host).await {
260251
error!("Error creating FlightSQL client: {:?}", e);
261252
if let Err(e) = _event_tx.send(AppEvent::FlightSQLFailedToConnect) {
262253
error!("Error sending FlightSQLFailedToConnect message: {e}");

‎tests/server_cases/flightsql.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,18 @@ pub async fn test_flightsql_custom_host() {
3232
// Drop the listener so that the port becomes available for the test server.
3333
drop(listener);
3434

35-
// Create the custom host using the chosen port.
36-
let custom_host = format!("127.0.0.1:{}", port);
35+
let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind to a random port");
36+
let metrics_port = listener.local_addr().unwrap().port();
37+
// Drop the listener so that the port becomes available for the test server.
38+
drop(listener);
3739

3840
// Seems like we need to assign variable for server not to drop and kill prematurely
3941
let _server = TokioCommand::new(bin)
4042
.arg("serve-flightsql")
41-
.arg("--host")
42-
.arg(custom_host)
43+
.arg("--port")
44+
.arg(format!("{port}"))
45+
.arg("--metrics-port")
46+
.arg(format!("{metrics_port}"))
4347
.kill_on_drop(true)
4448
.spawn()
4549
.expect("Failed to spawn the server");
@@ -53,7 +57,7 @@ pub async fn test_flightsql_custom_host() {
5357
.arg("SELECT 1")
5458
.arg("--flightsql")
5559
.arg("--host")
56-
.arg(format!("http://127.0.0.1:{}", port))
60+
.arg(format!("http://localhost:{}", port))
5761
.assert()
5862
.success();
5963

‎tests/server_cases/http.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,16 +29,21 @@ pub async fn test_http_custom_host() {
2929
let port = listener.local_addr().unwrap().port();
3030
// Drop the listener so that the port becomes available for the test server.
3131
drop(listener);
32+
let custom_host = format!("localhost:{port}");
3233

33-
// Create the custom host using the chosen port.
34-
let custom_host = format!("127.0.0.1:{}", port);
34+
let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind to a random port");
35+
let metrics_port = listener.local_addr().unwrap().port();
36+
// Drop the listener so that the port becomes available for the test server.
37+
drop(listener);
3538

3639
// Seems like we need to assign variable for server not to drop and kill prematurely
3740
let _server = TokioCommand::new(bin)
3841
.env("RUST_LOG", "off")
3942
.arg("serve-http")
40-
.arg("--host")
41-
.arg(&custom_host)
43+
.arg("--port")
44+
.arg(format!("{port}"))
45+
.arg("--metrics-port")
46+
.arg(format!("{metrics_port}"))
4247
.kill_on_drop(true)
4348
.spawn()
4449
.expect("Failed to spawn server");

0 commit comments

Comments
 (0)