Fix schema migration proxy (#1572)

fix shared schema request proxy
tursodatabase · Jul 18, 2024 · 188644b · 188644b
1 parent dfeb764
commit 188644b
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 12 deletions.
diff --git a/libsql-server/src/database/mod.rs b/libsql-server/src/database/mod.rs
@@ -2,9 +2,10 @@ use std::fmt;
 use std::sync::Arc;
 
 use bottomless::replicator::Replicator;
+use tokio::sync::watch;
 
 use crate::connection::{MakeConnection, RequestContext};
-use crate::replication::ReplicationLogger;
+use crate::replication::{FrameNo, ReplicationLogger};
 
 pub use self::primary::{PrimaryConnection, PrimaryConnectionMaker, PrimaryDatabase};
 pub use self::replica::{ReplicaConnection, ReplicaDatabase};
@@ -74,7 +75,7 @@ impl crate::connection::Connection for Connection {
         pgm: crate::connection::program::Program,
         ctx: RequestContext,
         response_builder: B,
-        replication_index: Option<crate::replication::FrameNo>,
+        replication_index: Option<FrameNo>,
     ) -> crate::Result<B> {
         match self {
             Connection::Primary(conn) => {
@@ -96,7 +97,7 @@ impl crate::connection::Connection for Connection {
         &self,
         sql: String,
         ctx: RequestContext,
-        replication_index: Option<crate::replication::FrameNo>,
+        replication_index: Option<FrameNo>,
     ) -> crate::Result<crate::Result<crate::connection::program::DescribeResponse>> {
         match self {
             Connection::Primary(conn) => conn.describe(sql, ctx, replication_index).await,
@@ -187,6 +188,26 @@ impl Database {
         }
     }
 
+    pub fn notifier(&self) -> Option<watch::Receiver<Option<FrameNo>>> {
+        match self {
+            Database::Primary(p) => Some(
+                p.wal_wrapper
+                    .wrapper()
+                    .logger()
+                    .new_frame_notifier
+                    .subscribe(),
+            ),
+            Database::Replica(_) => None,
+            Database::Schema(s) => Some(
+                s.wal_wrapper
+                    .wrapper()
+                    .logger()
+                    .new_frame_notifier
+                    .subscribe(),
+            ),
+        }
+    }
+
     pub fn as_primary(&self) -> Option<&PrimaryDatabase> {
         if let Self::Primary(v) = self {
             Some(v)

diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
@@ -581,15 +581,7 @@ impl Proxy for ProxyService {
             .namespaces
             .with(ctx.namespace().clone(), |ns| {
                 let connection_maker = ns.db.connection_maker();
-                let notifier = ns
-                    .db
-                    .as_primary()
-                    .expect("invalid call to stream_exec: not a primary")
-                    .wal_wrapper
-                    .wrapper()
-                    .logger()
-                    .new_frame_notifier
-                    .subscribe();
+                let notifier = ns.db.notifier();
                 (connection_maker, notifier)
             })
             .await