[fence 5/9] libsql-server: migration target lifecycle and import capability

libsql-server/src/admin_shell.rs

@@ -343,4 +343,95 @@ mod fence_tests {
         );
         assert_eq!(s.fence.read_lease_counts().total(), 0);
     }
+
+    /// Section 17 row 12: a quarantined target is written only through its operation's import
+    /// capability. The admin shell, which reaches the namespace with admin authority and runs raw
+    /// SQL, can neither read nor write it: every query is refused by the read admission with the
+    /// quarantine code, and a raw write that skipped the admission is still refused at the WAL.
+    /// The import session keeps working, and nothing the shell sent changed the data.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn admin_shell_cannot_write_quarantined() {
+        use crate::namespace::fence::target::tests::{create, create_request, OP as TARGET_OP};
+        use crate::namespace::open_test_store as open_store;
+
+        let dir = tempfile::tempdir().unwrap();
+        let store = open_store(dir.path()).await;
+        create(&store, create_request("tgt", 1))
+            .await
+            .unwrap()
+            .unwrap();
+        let mut session = store
+            .open_import_session("tgt".into(), TARGET_OP, 1)
+            .await
+            .unwrap();
+        session
+            .with_raw(|c| c.execute_batch("create table t (x); insert into t values (1)"))
+            .await
+            .unwrap()
+            .unwrap();
+
+        let shell = AdminShell::new(store.clone());
+        let sql = [
+            "insert into t values (2)",
+            "delete from t",
+            "create table u (y)",
+            "pragma user_version = 7",
+            "begin immediate",
+            "select count(*) from t",
+        ];
+        let queries = tokio_stream::iter(sql.map(|q| Ok(rpc::Query { query: q.into() })));
+        let responses: Vec<_> = shell
+            .with_namespace(Bytes::from_static(b"tgt"), queries)
+            .await
+            .unwrap()
+            .collect()
+            .await;
+        assert_eq!(responses.len(), sql.len());
+        for (q, resp) in sql.iter().zip(&responses) {
+            let resp = resp.as_ref().unwrap();
+            assert!(
+                error(resp).starts_with("MIGRATION_TARGET_QUARANTINED"),
+                "{q}: {}",
+                error(resp)
+            );
+        }
+
+        // Without the shell's read admission, the raw write is refused by the WAL gate: the
+        // connection holds no capability.
+        let (fence, maker) = store
+            .with("tgt".into(), |ns| {
+                (ns.fence().clone(), ns.db.connection_maker())
+            })
+            .await
+            .unwrap();
+        let conn = maker.create().await.unwrap();
+        for q in ["insert into t values (3)", "create table v (z)"] {
+            let resp = conn.with_raw(|c| run_one(c, q.into())).unwrap();
+            assert!(error(&resp).contains("authoriz"), "{q}: {}", error(&resp));
+        }
+        assert_eq!(fence.read_lease_counts().total(), 0);
+
+        // The capability still writes, and it sees only its own rows.
+        let rows: i64 = session
+            .with_raw(|c| {
+                c.execute("insert into t values (4)", ())?;
+                c.query_row("select count(*) from t", (), |r| r.get(0))
+            })
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(rows, 2);
+        let tables: i64 = session
+            .with_raw(|c| {
+                c.query_row(
+                    "select count(*) from sqlite_schema where type = 'table'",
+                    (),
+                    |r| r.get(0),
+                )
+            })
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(tables, 1);
+    }
 }

libsql-server/src/connection/legacy.rs

@@ -14,6 +14,7 @@ use tokio::time::Duration;
 use crate::error::Error;
 use crate::metrics::DESCRIBE_COUNT;
 use crate::namespace::broadcasters::BroadcasterHandle;
+use crate::namespace::fence::capability::MigrationCapability;
 use crate::namespace::fence::controller::{FenceConnState, FenceController};
 use crate::namespace::fence::state::OperationClass;
 use crate::namespace::meta_store::MetaStoreHandle;
@@ -143,6 +144,33 @@ where
 
     #[tracing::instrument(skip(self))]
     pub(super) async fn make_connection(&self) -> Result<LegacyConnection<W>> {
+        self.make_connection_with(FenceConnState::new(
+            self.fence.clone(),
+            OperationClass::NormalWrite,
+        ))
+        .await
+    }
+
+    /// Open a connection that works under `capability` (an import or validation session,
+    /// `docs/NAMESPACE_FENCE.md` section 11). It shares the maker's write slot, WAL and
+    /// replication log with every other connection, and is admitted as the capability's class
+    /// only while the capability is valid. It is not counted by the connection throttle: it is
+    /// operation-owned work, and the fence, not the throttle, bounds how much of it runs.
+    pub(crate) async fn make_capability_connection(
+        &self,
+        capability: MigrationCapability,
+    ) -> Result<LegacyConnection<W>> {
+        self.make_connection_with(FenceConnState::with_capability(
+            self.fence.clone(),
+            capability,
+        ))
+        .await
+    }
+
+    async fn make_connection_with(
+        &self,
+        fence: Arc<FenceConnState>,
+    ) -> Result<LegacyConnection<W>> {
         LegacyConnection::new(
             self.db_path.clone(),
             self.extensions.clone(),
@@ -161,7 +189,7 @@ where
             self.resolve_attach_path.clone(),
             self.connection_manager.clone(),
             self.make_wal_manager.clone(),
-            FenceConnState::new(self.fence.clone(), OperationClass::NormalWrite),
+            fence,
         )
         .await
     }
@@ -213,6 +241,13 @@ impl LegacyConnection<libsql_sys::wal::wrapper::PassthroughWalWrapper> {
     }
 }
 
+impl<T> LegacyConnection<T> {
+    /// The fence state shared by this connection's WAL wrapper and `CoreConnection`.
+    pub(crate) fn fence_state(&self) -> &Arc<FenceConnState> {
+        &self.fence
+    }
+}
+
 impl<T> Clone for LegacyConnection<T> {
     fn clone(&self) -> Self {
         Self {

libsql-server/src/connection/mod.rs

@@ -286,6 +286,11 @@ pub struct MakeThrottledConnection<F> {
 }
 
 impl<F> MakeThrottledConnection<F> {
+    /// The connection maker this one throttles.
+    pub(crate) fn inner(&self) -> &F {
+        &self.connection_maker
+    }
+
     fn new(
         semaphore: Arc<Semaphore>,
         connection_maker: F,

libsql-server/src/namespace/dump_import/buffered.rs

@@ -13,19 +13,21 @@ use sqlite3_parser::lexer::sql::Parser;
 use tokio::io::AsyncReadExt;
 use tokio_util::io::StreamReader;
 
-use crate::connection::Connection as _;
-use crate::database::PrimaryConnection;
+use crate::connection::Connection;
 use crate::error::LoadDumpError;
 
 use super::{map_exec_error, map_parse_error, DumpImportStats};
 
-pub(super) async fn load_dump_buffered<S>(
+pub(super) async fn load_dump_buffered<S, C>(
     dump: S,
-    conn: PrimaryConnection,
+    conn: C,
 ) -> crate::Result<DumpImportStats, LoadDumpError>
 where
     S: Stream<Item = std::io::Result<Bytes>> + Unpin,
+    C: Connection,
 {
+    // Shared with the blocking task that runs each statement.
+    let conn = std::sync::Arc::new(conn);
     let mut stats = DumpImportStats::default();
     let mut reader = tokio::io::BufReader::new(StreamReader::new(dump));
     let mut dump_content = String::new();

libsql-server/src/namespace/dump_import/mod.rs

@@ -18,7 +18,7 @@ use std::time::{Duration, Instant};
 
 use sqlite3_parser::lexer::sql::ParserError;
 
-use crate::database::PrimaryConnection;
+use crate::connection::Connection;
 use crate::error::LoadDumpError;
 use crate::namespace::{DumpStream, NamespaceName};
 
@@ -163,10 +163,10 @@ impl From<LoadDumpError> for ImportFailure {
 /// Load `stream` into the fresh database behind `conn` using the requested importer.
 ///
 /// On error the dump transaction has been rolled back (or was never started).
-pub(crate) async fn load_dump(
+pub(crate) async fn load_dump<C: Connection>(
     kind: DumpImporterKind,
     stream: DumpStream,
-    conn: PrimaryConnection,
+    conn: C,
     cfg: &DumpImportConfig,
     namespace: &NamespaceName,
 ) -> Result<DumpImportStats, LoadDumpError> {

libsql-server/src/namespace/dump_import/streaming.rs

@@ -2,7 +2,7 @@
 //!
 //! ```text
 //!          async (tokio runtime)                                blocking thread (BLOCKING_RT)
-//! DumpStream ──chunks──▶ StatementFramer ──frames──▶ mpsc + byte budget ──▶ run_executor ──▶ PrimaryConnection
+//! DumpStream ──chunks──▶ StatementFramer ──frames──▶ mpsc + byte budget ──▶ run_executor ──▶ Connection
 //! ```
 //!
 //! The reader task drives the stream, frames statements and hands them to a single executor
@@ -21,8 +21,7 @@ use sqlite3_parser::ast::{Cmd, Stmt};
 use sqlite3_parser::lexer::sql::Parser;
 use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore};
 
-use crate::connection::Connection as _;
-use crate::database::PrimaryConnection;
+use crate::connection::Connection;
 use crate::error::LoadDumpError;
 use crate::namespace::DumpStream;
 use crate::BLOCKING_RT;
@@ -68,9 +67,9 @@ impl From<FrameError> for LoadDumpError {
     }
 }
 
-pub(super) async fn load_dump_streaming(
+pub(super) async fn load_dump_streaming<C: Connection>(
     mut stream: DumpStream,
-    conn: PrimaryConnection,
+    conn: C,
     cfg: &DumpImportConfig,
 ) -> Result<DumpImportStats, ImportFailure> {
     // `DumpImportConfig::validate` runs for CLI-built configs only; never panic on a
@@ -157,8 +156,8 @@ async fn send_frame(
 /// How often the executor logs progress while an import is running.
 const PROGRESS_LOG_INTERVAL: Duration = Duration::from_secs(10);
 
-struct Executor {
-    conn: PrimaryConnection,
+struct Executor<C> {
+    conn: C,
     n_stmt: u64,
     skipped_wasm_table: bool,
     stats: DumpImportStats,
@@ -173,8 +172,8 @@ struct Executor {
 ///
 /// Returns the outcome together with the statistics accumulated so far, so a failure can report
 /// how far the import got.
-fn run_executor(
-    conn: PrimaryConnection,
+fn run_executor<C: Connection>(
+    conn: C,
     mut rx: mpsc::Receiver<Msg>,
 ) -> (Result<(), LoadDumpError>, DumpImportStats) {
     let now = Instant::now();
@@ -204,7 +203,7 @@ fn run_executor(
     (res, std::mem::take(&mut ex.stats))
 }
 
-impl Executor {
+impl<C: Connection> Executor<C> {
     fn run(&mut self, rx: &mut mpsc::Receiver<Msg>) -> Result<(), LoadDumpError> {
         loop {
             match rx.blocking_recv() {