Skip to content

Commit 91611e7

Browse files
jameswritescodejameswritescode
committed
Add support for namespace to the remote connection builder
1 parent 9241b00 commit 91611e7

File tree

6 files changed

+94
-10
lines changed

6 files changed

+94
-10
lines changed

‎libsql-server/src/http/user/db_factory.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ pub fn namespace_from_headers(
5050

5151
if let Some(from_metadata) = headers.get(NAMESPACE_METADATA_KEY) {
5252
try_namespace_from_metadata(from_metadata)
53+
} else if let Some(from_ns_header) = headers.get("x-namespace") {
54+
try_namespace_from_header(from_ns_header)
5355
} else if let Some(from_host) = headers.get("host") {
5456
try_namespace_from_host(from_host, disable_default_namespace)
5557
} else if !disable_default_namespace {
@@ -59,6 +61,11 @@ pub fn namespace_from_headers(
5961
}
6062
}
6163

64+
fn try_namespace_from_header(header: &axum::http::HeaderValue) -> Result<NamespaceName, Error> {
65+
NamespaceName::from_bytes(header.as_bytes().to_vec().into())
66+
.map_err(|_| Error::InvalidNamespace)
67+
}
68+
6269
fn try_namespace_from_host(
6370
from_host: &axum::http::HeaderValue,
6471
disable_default_namespace: bool,

‎libsql-server/tests/embedded_replica/mod.rs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1696,3 +1696,43 @@ fn schema_db() {
16961696

16971697
sim.run().unwrap();
16981698
}
1699+
1700+
#[test]
1701+
fn remote_namespace_header_support() {
1702+
let tmp_host = tempdir().unwrap();
1703+
let tmp_host_path = tmp_host.path().to_owned();
1704+
1705+
let mut sim = Builder::new()
1706+
.simulation_duration(Duration::from_secs(1000))
1707+
.build();
1708+
1709+
make_primary(&mut sim, tmp_host_path.clone());
1710+
1711+
sim.client("client", async move {
1712+
let client = Client::new();
1713+
1714+
client
1715+
.post("http://primary:9090/v1/namespaces/foo/create", json!({}))
1716+
.await?;
1717+
1718+
let db_url = "http://primary:8080";
1719+
1720+
let remote = libsql::Builder::new_remote(db_url.to_string(), String::new())
1721+
.namespace("foo")
1722+
.connector(TurmoilConnector)
1723+
.build()
1724+
.await
1725+
.unwrap();
1726+
1727+
let conn = remote.connect().unwrap();
1728+
1729+
conn.execute("CREATE TABLE user (id INTEGER NOT NULL PRIMARY KEY)", ())
1730+
.await?;
1731+
1732+
conn.execute("INSERT into user(id) values (1);", ()).await?;
1733+
1734+
Ok(())
1735+
});
1736+
1737+
sim.run().unwrap();
1738+
}

‎libsql/src/database.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ enum DbType {
8888
auth_token: String,
8989
connector: crate::util::ConnectorService,
9090
version: Option<String>,
91+
namespace: Option<String>,
9192
},
9293
}
9394

@@ -230,6 +231,7 @@ cfg_replication! {
230231
OpenFlags::default(),
231232
encryption_config.clone(),
232233
None,
234+
None,
233235
).await?;
234236

235237
Ok(Database {
@@ -514,6 +516,7 @@ cfg_remote! {
514516
auth_token: auth_token.into(),
515517
connector: crate::util::ConnectorService::new(svc),
516518
version,
519+
namespace: None,
517520
},
518521
max_write_replication_index: Default::default(),
519522
})
@@ -650,13 +653,15 @@ impl Database {
650653
auth_token,
651654
connector,
652655
version,
656+
namespace,
653657
} => {
654658
let conn = std::sync::Arc::new(
655659
crate::hrana::connection::HttpConnection::new_with_connector(
656660
url,
657661
auth_token,
658662
connector.clone(),
659663
version.as_ref().map(|s| s.as_str()),
664+
namespace.as_ref().map(|s| s.as_str()),
660665
),
661666
);
662667

‎libsql/src/database/builder.rs

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,12 +59,12 @@ impl Builder<()> {
5959
auth_token,
6060
connector: None,
6161
version: None,
62+
namespace: None,
6263
},
6364
encryption_config: None,
6465
read_your_writes: true,
6566
sync_interval: None,
6667
http_request_callback: None,
67-
namespace: None
6868
},
6969
}
7070
}
@@ -99,6 +99,7 @@ impl Builder<()> {
9999
auth_token,
100100
connector: None,
101101
version: None,
102+
namespace: None,
102103
},
103104
connector:None,
104105
},
@@ -115,6 +116,7 @@ impl Builder<()> {
115116
auth_token,
116117
connector: None,
117118
version: None,
119+
namespace: None,
118120
},
119121
}
120122
}
@@ -128,6 +130,7 @@ cfg_replication_or_remote_or_sync! {
128130
auth_token: String,
129131
connector: Option<crate::util::ConnectorService>,
130132
version: Option<String>,
133+
namespace: Option<String>,
131134
}
132135
}
133136

@@ -195,7 +198,6 @@ cfg_replication! {
195198
read_your_writes: bool,
196199
sync_interval: Option<std::time::Duration>,
197200
http_request_callback: Option<crate::util::HttpRequestCallback>,
198-
namespace: Option<String>,
199201
}
200202

201203
/// Local replica configuration type in [`Builder`].
@@ -260,7 +262,7 @@ cfg_replication! {
260262
/// Set the namespace that will be communicated to remote replica in the http header.
261263
pub fn namespace(mut self, namespace: impl Into<String>) -> Builder<RemoteReplica>
262264
{
263-
self.inner.namespace = Some(namespace.into());
265+
self.inner.remote.namespace = Some(namespace.into());
264266
self
265267
}
266268

@@ -280,12 +282,12 @@ cfg_replication! {
280282
auth_token,
281283
connector,
282284
version,
285+
namespace,
283286
},
284287
encryption_config,
285288
read_your_writes,
286289
sync_interval,
287290
http_request_callback,
288-
namespace
289291
} = self.inner;
290292

291293
let connector = if let Some(connector) = connector {
@@ -357,6 +359,7 @@ cfg_replication! {
357359
auth_token,
358360
connector,
359361
version,
362+
namespace,
360363
}) = remote
361364
{
362365
let connector = if let Some(connector) = connector {
@@ -381,6 +384,7 @@ cfg_replication! {
381384
flags,
382385
encryption_config.clone(),
383386
http_request_callback,
387+
namespace,
384388
)
385389
.await?
386390
} else {
@@ -434,6 +438,7 @@ cfg_sync! {
434438
auth_token,
435439
connector: _,
436440
version: _,
441+
namespace: _,
437442
},
438443
connector,
439444
} = self.inner;
@@ -490,13 +495,21 @@ cfg_remote! {
490495
self
491496
}
492497

498+
/// Set the namespace that will be communicated to the remote in the http header.
499+
pub fn namespace(mut self, namespace: impl Into<String>) -> Builder<Remote>
500+
{
501+
self.inner.namespace = Some(namespace.into());
502+
self
503+
}
504+
493505
/// Build the remote database client.
494506
pub async fn build(self) -> Result<Database> {
495507
let Remote {
496508
url,
497509
auth_token,
498510
connector,
499511
version,
512+
namespace,
500513
} = self.inner;
501514

502515
let connector = if let Some(connector) = connector {
@@ -518,6 +531,7 @@ cfg_remote! {
518531
auth_token,
519532
connector,
520533
version,
534+
namespace,
521535
},
522536
max_write_replication_index: Default::default(),
523537
})

‎libsql/src/hrana/hyper.rs

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,27 @@ pub type ByteStream = Box<dyn Stream<Item = std::io::Result<Bytes>> + Send + Syn
2525
pub struct HttpSender {
2626
inner: hyper::Client<ConnectorService, hyper::Body>,
2727
version: HeaderValue,
28+
namespace: Option<HeaderValue>,
2829
}
2930

3031
impl HttpSender {
31-
pub fn new(connector: ConnectorService, version: Option<&str>) -> Self {
32+
pub fn new(
33+
connector: ConnectorService,
34+
version: Option<&str>,
35+
namespace: Option<&str>,
36+
) -> Self {
3237
let ver = version.unwrap_or(env!("CARGO_PKG_VERSION"));
3338

3439
let version = HeaderValue::try_from(format!("libsql-remote-{ver}")).unwrap();
40+
let namespace = namespace.map(|v| HeaderValue::try_from(v).unwrap());
3541

3642
let inner = hyper::Client::builder().build(connector);
3743

38-
Self { inner, version }
44+
Self {
45+
inner,
46+
version,
47+
namespace,
48+
}
3949
}
4050

4151
async fn send(
@@ -44,9 +54,15 @@ impl HttpSender {
4454
auth: Arc<str>,
4555
body: String,
4656
) -> Result<super::HttpBody<ByteStream>> {
47-
let req = hyper::Request::post(url.as_ref())
57+
let mut req_builder = hyper::Request::post(url.as_ref())
4858
.header(AUTHORIZATION, auth.as_ref())
49-
.header("x-libsql-client-version", self.version.clone())
59+
.header("x-libsql-client-version", self.version.clone());
60+
61+
if let Some(namespace) = self.namespace {
62+
req_builder = req_builder.header("x-namespace", namespace);
63+
}
64+
65+
let req = req_builder
5066
.body(hyper::Body::from(body))
5167
.map_err(|err| HranaError::Http(format!("{:?}", err)))?;
5268

@@ -107,8 +123,9 @@ impl HttpConnection<HttpSender> {
107123
token: impl Into<String>,
108124
connector: ConnectorService,
109125
version: Option<&str>,
126+
namespace: Option<&str>,
110127
) -> Self {
111-
let inner = HttpSender::new(connector, version);
128+
let inner = HttpSender::new(connector, version, namespace);
112129
Self::new(url.into(), token.into(), inner)
113130
}
114131
}

‎libsql/src/local/database.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,7 @@ impl Database {
190190
flags: OpenFlags,
191191
encryption_config: Option<EncryptionConfig>,
192192
http_request_callback: Option<crate::util::HttpRequestCallback>,
193+
namespace: Option<String>,
193194
) -> Result<Database> {
194195
use std::path::PathBuf;
195196

@@ -208,7 +209,7 @@ impl Database {
208209
auth_token,
209210
version.as_deref(),
210211
http_request_callback,
211-
None,
212+
namespace,
212213
)
213214
.map_err(|e| crate::Error::Replication(e.into()))?;
214215

0 commit comments

Comments
 (0)