Add support for namespace to the remote connection builder

Author: James Newton

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

@@ -50,6 +50,8 @@ pub fn namespace_from_headers(
 
     if let Some(from_metadata) = headers.get(NAMESPACE_METADATA_KEY) {
         try_namespace_from_metadata(from_metadata)
+    } else if let Some(from_ns_header) = headers.get("x-namespace") {
+        try_namespace_from_header(from_ns_header)
     } else if let Some(from_host) = headers.get("host") {
         try_namespace_from_host(from_host, disable_default_namespace)
     } else if !disable_default_namespace {
@@ -59,6 +61,11 @@ pub fn namespace_from_headers(
     }
 }
 
+fn try_namespace_from_header(header: &axum::http::HeaderValue) -> Result<NamespaceName, Error> {
+    NamespaceName::from_bytes(header.as_bytes().to_vec().into())
+        .map_err(|_| Error::InvalidNamespace)
+}
+
 fn try_namespace_from_host(
     from_host: &axum::http::HeaderValue,
     disable_default_namespace: bool,

libsql-server/tests/embedded_replica/mod.rs

@@ -1696,3 +1696,103 @@ fn schema_db() {
 
     sim.run().unwrap();
 }
+
+#[test]
+fn remote_namespace_header_support() {
+    let tmp_host = tempdir().unwrap();
+    let tmp_host_path = tmp_host.path().to_owned();
+
+    let mut sim = Builder::new()
+        .simulation_duration(Duration::from_secs(1000))
+        .build();
+
+    make_primary(&mut sim, tmp_host_path.clone());
+
+    sim.client("client", async move {
+        let client = Client::new();
+
+        client
+            .post("http://primary:9090/v1/namespaces/foo/create", json!({}))
+            .await?;
+
+        let db_url = "http://primary:8080";
+
+        let remote = libsql::Builder::new_remote(db_url.to_string(), String::new())
+            .namespace("foo")
+            .connector(TurmoilConnector)
+            .build()
+            .await
+            .unwrap();
+
+        let conn = remote.connect().unwrap();
+
+        conn.execute("CREATE TABLE user (id INTEGER NOT NULL PRIMARY KEY)", ())
+            .await?;
+
+        conn.execute("INSERT into user(id) values (1);", ()).await?;
+
+        Ok(())
+    });
+
+    sim.run().unwrap();
+}
+
+#[test]
+fn remote_replica_namespace_header_support() {
+    let tmp_host = tempdir().unwrap();
+    let tmp_host_path = tmp_host.path().to_owned();
+
+    let tmp_embedded = tempdir().unwrap();
+    let tmp_embedded_path = tmp_embedded.path().to_owned();
+
+    let mut sim = Builder::new()
+        .simulation_duration(Duration::from_secs(1000))
+        .build();
+
+    make_primary(&mut sim, tmp_host_path.clone());
+
+    sim.client("client", async move {
+        let client = Client::new();
+
+        client
+            .post("http://primary:9090/v1/namespaces/foo/create", json!({}))
+            .await?;
+
+        let db_url = "http://primary:8080";
+
+        let remote = libsql::Builder::new_remote(db_url.to_string(), String::new())
+            .namespace("foo")
+            .connector(TurmoilConnector)
+            .build()
+            .await
+            .unwrap();
+
+        let conn = remote.connect().unwrap();
+
+        conn.execute("CREATE TABLE user (id INTEGER NOT NULL PRIMARY KEY)", ())
+            .await?;
+
+        conn.execute("INSERT into user(id) values (1);", ()).await?;
+
+        let path = tmp_embedded_path.join("embedded");
+
+        let remote_replica = libsql::Builder::new_remote_replica(
+            path.to_str().unwrap(),
+            db_url.to_string(),
+            String::new(),
+        )
+        .namespace("foo")
+        .connector(TurmoilConnector)
+        .build()
+        .await
+        .unwrap();
+
+        let rep = remote_replica.sync().await.unwrap();
+        assert_eq!(rep.frame_no(), Some(2));
+        assert_eq!(rep.frames_synced(), 3);
+
+        Ok(())
+    });
+
+    sim.run().unwrap();
+}

libsql/src/database.rs

@@ -107,6 +107,7 @@ enum DbType {
         auth_token: String,
         connector: crate::util::ConnectorService,
         version: Option<String>,
+        namespace: Option<String>,
     },
 }
 
@@ -250,6 +251,7 @@ cfg_replication! {
                 OpenFlags::default(),
                 encryption_config.clone(),
                 None,
+                None,
             ).await?;
 
             Ok(Database {
@@ -541,6 +543,7 @@ cfg_remote! {
                     auth_token: auth_token.into(),
                     connector: crate::util::ConnectorService::new(svc),
                     version,
+                    namespace: None,
                 },
                 max_write_replication_index: Default::default(),
             })
@@ -704,6 +707,7 @@ impl Database {
                             auth_token.clone(),
                             connector.clone(),
                             None,
+                            None,
                         ),
                         read_your_writes: *read_your_writes,
                         context: db.sync_ctx.clone().unwrap(),
@@ -724,13 +728,15 @@ impl Database {
                 auth_token,
                 connector,
                 version,
+                namespace,
             } => {
                 let conn = std::sync::Arc::new(
                     crate::hrana::connection::HttpConnection::new_with_connector(
                         url,
                         auth_token,
                         connector.clone(),
                         version.as_ref().map(|s| s.as_str()),
+                        namespace.as_ref().map(|s| s.as_str()),
                     ),
                 );
 

libsql/src/database/builder.rs

@@ -59,12 +59,12 @@ impl Builder<()> {
                         auth_token,
                         connector: None,
                         version: None,
+                        namespace: None,
                     },
                     encryption_config: None,
                     read_your_writes: true,
                     sync_interval: None,
                     http_request_callback: None,
-                    namespace: None,
                     skip_safety_assert: false,
                     #[cfg(feature = "sync")]
                     sync_protocol: Default::default(),
@@ -102,6 +102,7 @@ impl Builder<()> {
                         auth_token,
                         connector: None,
                         version: None,
+                        namespace: None,
                     },
                     connector: None,
                     read_your_writes: true,
@@ -122,6 +123,7 @@ impl Builder<()> {
                     auth_token,
                     connector: None,
                     version: None,
+                    namespace: None,
                 },
             }
         }
@@ -135,6 +137,7 @@ cfg_replication_or_remote_or_sync! {
         auth_token: String,
         connector: Option<crate::util::ConnectorService>,
         version: Option<String>,
+        namespace: Option<String>,
     }
 }
 
@@ -223,7 +226,6 @@ cfg_replication! {
         read_your_writes: bool,
         sync_interval: Option<std::time::Duration>,
         http_request_callback: Option<crate::util::HttpRequestCallback>,
-        namespace: Option<String>,
         skip_safety_assert: bool,
         #[cfg(feature = "sync")]
         sync_protocol: super::SyncProtocol,
@@ -300,7 +302,7 @@ cfg_replication! {
         /// Set the namespace that will be communicated to remote replica in the http header.
         pub fn namespace(mut self, namespace: impl Into<String>) -> Builder<RemoteReplica>
         {
-            self.inner.namespace = Some(namespace.into());
+            self.inner.remote.namespace = Some(namespace.into());
             self
         }
 
@@ -334,12 +336,12 @@ cfg_replication! {
                         auth_token,
                         connector,
                         version,
+                        namespace,
                     },
                 encryption_config,
                 read_your_writes,
                 sync_interval,
                 http_request_callback,
-                namespace,
                 skip_safety_assert,
                 #[cfg(feature = "sync")]
                 sync_protocol,
@@ -500,6 +502,7 @@ cfg_replication! {
                 auth_token,
                 connector,
                 version,
+                namespace,
             }) = remote
             {
                 let connector = if let Some(connector) = connector {
@@ -524,6 +527,7 @@ cfg_replication! {
                     flags,
                     encryption_config.clone(),
                     http_request_callback,
+                    namespace,
                 )
                 .await?
             } else {
@@ -606,6 +610,7 @@ cfg_sync! {
                         auth_token,
                         connector: _,
                         version: _,
+                        namespace: _,
                     },
                 connector,
                 remote_writes,
@@ -730,13 +735,21 @@ cfg_remote! {
             self
         }
 
+        /// Set the namespace that will be communicated to the remote in the http header.
+        pub fn namespace(mut self, namespace: impl Into<String>) -> Builder<Remote>
+        {
+            self.inner.namespace = Some(namespace.into());
+            self
+        }
+
         /// Build the remote database client.
         pub async fn build(self) -> Result<Database> {
             let Remote {
                 url,
                 auth_token,
                 connector,
                 version,
+                namespace,
             } = self.inner;
 
             let connector = if let Some(connector) = connector {
@@ -758,6 +771,7 @@ cfg_remote! {
                     auth_token,
                     connector,
                     version,
+                    namespace,
                 },
                 max_write_replication_index: Default::default(),
             })

libsql/src/hrana/hyper.rs

@@ -26,17 +26,27 @@ pub type ByteStream = Box<dyn Stream<Item = std::io::Result<Bytes>> + Send + Syn
 pub struct HttpSender {
     inner: hyper::Client<ConnectorService, hyper::Body>,
     version: HeaderValue,
+    namespace: Option<HeaderValue>,
 }
 
 impl HttpSender {
-    pub fn new(connector: ConnectorService, version: Option<&str>) -> Self {
+    pub fn new(
+        connector: ConnectorService,
+        version: Option<&str>,
+        namespace: Option<&str>,
+    ) -> Self {
         let ver = version.unwrap_or(env!("CARGO_PKG_VERSION"));
 
         let version = HeaderValue::try_from(format!("libsql-remote-{ver}")).unwrap();
+        let namespace = namespace.map(|v| HeaderValue::try_from(v).unwrap());
 
         let inner = hyper::Client::builder().build(connector);
 
-        Self { inner, version }
+        Self {
+            inner,
+            version,
+            namespace,
+        }
     }
 
     async fn send(
@@ -45,9 +55,15 @@ impl HttpSender {
         auth: Arc<str>,
         body: String,
     ) -> Result<super::HttpBody<ByteStream>> {
-        let req = hyper::Request::post(url.as_ref())
+        let mut req_builder = hyper::Request::post(url.as_ref())
             .header(AUTHORIZATION, auth.as_ref())
-            .header("x-libsql-client-version", self.version.clone())
+            .header("x-libsql-client-version", self.version.clone());
+
+        if let Some(namespace) = self.namespace {
+            req_builder = req_builder.header("x-namespace", namespace);
+        }
+
+        let req = req_builder
             .body(hyper::Body::from(body))
             .map_err(|err| HranaError::Http(format!("{:?}", err)))?;
 
@@ -109,8 +125,9 @@ impl HttpConnection<HttpSender> {
         token: impl Into<String>,
         connector: ConnectorService,
         version: Option<&str>,
+        namespace: Option<&str>,
     ) -> Self {
-        let inner = HttpSender::new(connector, version);
+        let inner = HttpSender::new(connector, version, namespace);
         Self::new(url.into(), token.into(), inner)
     }
 }

libsql/src/local/database.rs

@@ -266,6 +266,7 @@ impl Database {
         flags: OpenFlags,
         encryption_config: Option<EncryptionConfig>,
         http_request_callback: Option<crate::util::HttpRequestCallback>,
+        namespace: Option<String>,
     ) -> Result<Database> {
         use std::path::PathBuf;
 
@@ -284,7 +285,7 @@ impl Database {
             auth_token,
             version.as_deref(),
             http_request_callback,
-            None,
+            namespace,
         )
         .map_err(|e| crate::Error::Replication(e.into()))?;