fix test

Author: joshua-spacetime

crates/cli/src/subcommands/subscribe.rs

@@ -6,7 +6,7 @@ use reqwest::Url;
 use serde_json::Value;
 use spacetimedb_client_api_messages::websocket::{self as ws, JsonFormat};
 use spacetimedb_data_structures::map::HashMap;
-use spacetimedb_lib::db::raw_def::v9::{RawMiscModuleExportV9, RawModuleDefV9};
+use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9;
 use spacetimedb_lib::de::serde::{DeserializeWrapper, SeedWrapper};
 use spacetimedb_lib::ser::serde::SerializeWrapper;
 use std::io;
@@ -86,33 +86,11 @@ fn reformat_update<'a>(
     msg.tables
         .iter()
         .map(|upd| {
-            let product_type_ref = schema
-                .tables
-                .iter()
-                .find(|tbl| tbl.name == upd.table_name)
-                .map(|table_def| table_def.product_type_ref)
-                .or_else(|| {
-                    schema
-                        .misc_exports
-                        .iter()
-                        .filter_map(|misc_export| match misc_export {
-                            RawMiscModuleExportV9::View(view_def) => Some(view_def),
-                            _ => None,
-                        })
-                        .find(|view_def| view_def.name == upd.table_name)
-                        .map(|view_def| view_def.return_type.clone())
-                        .and_then(|return_type| {
-                            return_type
-                                .as_option()
-                                .map(|inner| inner.clone().into_ref())
-                                .or_else(|| return_type.as_array().map(|inner| inner.elem_ty.clone().into_ref()))
-                                .transpose()
-                                .expect("views must return an option or a vec")
-                        })
-                })
-                .context("table not found in schema")?;
-
-            let table_ty = schema.typespace.resolve(product_type_ref);
+            let table_ty = schema.typespace.resolve(
+                schema
+                    .type_ref_for_table_like(&upd.table_name)
+                    .context("table not found in schema")?,
+            );
 
             let reformat_row = |row: &str| -> anyhow::Result<Value> {
                 // TODO: can the following two calls be merged into a single call to reduce allocations?

crates/lib/src/db/raw_def/v9.rs

@@ -97,6 +97,50 @@ pub struct RawModuleDefV9 {
     pub row_level_security: Vec<RawRowLevelSecurityDefV9>,
 }
 
+impl RawModuleDefV9 {
+    /// Find a [`RawTableDefV9`] by name in this raw module def
+    fn find_table_def(&self, table_name: &str) -> Option<&RawTableDefV9> {
+        self.tables
+            .iter()
+            .find(|table_def| table_def.name.as_ref() == table_name)
+    }
+
+    /// Find a [`RawViewDefV9`] by name in this raw module def
+    fn find_view_def(&self, view_name: &str) -> Option<&RawViewDefV9> {
+        self.misc_exports.iter().find_map(|misc_export| match misc_export {
+            RawMiscModuleExportV9::View(view_def) if view_def.name.as_ref() == view_name => Some(view_def),
+            _ => None,
+        })
+    }
+
+    /// Find and return the product type ref for a table in this module def
+    fn type_ref_for_table(&self, table_name: &str) -> Option<AlgebraicTypeRef> {
+        self.find_table_def(table_name)
+            .map(|table_def| table_def.product_type_ref)
+    }
+
+    /// Find and return the product type ref for a view in this module def
+    fn type_ref_for_view(&self, view_name: &str) -> Option<AlgebraicTypeRef> {
+        self.find_view_def(view_name)
+            .map(|view_def| &view_def.return_type)
+            .and_then(|return_type| {
+                return_type
+                    .as_option()
+                    .and_then(|inner| inner.clone().into_ref().ok())
+                    .or_else(|| {
+                        return_type
+                            .as_array()
+                            .and_then(|inner| inner.elem_ty.clone().into_ref().ok())
+                    })
+            })
+    }
+
+    /// Find and return the product type ref for a table or view in this module def
+    pub fn type_ref_for_table_like(&self, name: &str) -> Option<AlgebraicTypeRef> {
+        self.type_ref_for_table(name).or_else(|| self.type_ref_for_view(name))
+    }
+}
+
 /// The definition of a database table.
 ///
 /// This struct holds information about the table, including its name, columns, indexes,

smoketests/tests/views.py

@@ -458,7 +458,7 @@ def test_subscribing_with_different_identities(self):
         self.call("insert_player", "Bob")
 
         # Subscribe to `my_player` as Bob
-        sub = self.subscribe("select * from my_player", n=1)
+        sub = self.subscribe("select * from my_player")
         events = sub()
 
         # Project out the identity field.