fix: stream source can't get the same tablemeta (#14935)

zhyass · web-flow · commit a10bd38bf3c1 · 2024-03-14T08:57:48.000Z
* fix: stream can't get the same source tablemeta

* add stateful test
diff --git a/src/query/catalog/src/catalog/interface.rs b/src/query/catalog/src/catalog/interface.rs
@@ -338,15 +338,13 @@ pub trait Catalog: DynClone + Send + Sync + Debug {
         unimplemented!()
     }
 
-    async fn stream_source_table(
-        &self,
-        _stream_desc: &str,
-        _tenant: &str,
-        _db_name: &str,
-        _source_table_name: &str,
-    ) -> Result<Arc<dyn Table>> {
+    fn get_stream_source_table(&self, _stream_desc: &str) -> Result<Option<Arc<dyn Table>>> {
         Err(ErrorCode::Unimplemented(
-            "'stream_source_table' not implemented",
+            "'get_stream_source_table' not implemented",
         ))
     }
+
+    fn cache_stream_source_table(&self, _stream: TableInfo, _source: TableInfo) {
+        unimplemented!()
+    }
 }
diff --git a/src/query/catalog/src/catalog/session_catalog.rs b/src/query/catalog/src/catalog/session_catalog.rs
@@ -289,14 +289,9 @@ impl Catalog for SessionCatalog {
                 } else {
                     let table = self.inner.get_table(tenant, db_name, table_name).await?;
                     if table.engine() == "STREAM" {
-                        let source_table = table
-                            .stream_source_table(Arc::new(self.clone()))
-                            .await?
-                            .get_table_info()
-                            .clone();
                         self.txn_mgr
                             .lock()
-                            .add_stream_table(table.get_table_info().clone(), source_table);
+                            .upsert_table_desc_to_id(table.get_table_info().clone());
                     }
                     Ok(table)
                 }
@@ -467,26 +462,25 @@ impl Catalog for SessionCatalog {
         self.inner.get_table_engines()
     }
 
-    async fn stream_source_table(
-        &self,
-        stream_desc: &str,
-        tenant: &str,
-        db_name: &str,
-        source_table_name: &str,
-    ) -> Result<Arc<dyn Table>> {
+    // Get stream source table from buffer by stream desc.
+    fn get_stream_source_table(&self, stream_desc: &str) -> Result<Option<Arc<dyn Table>>> {
         let is_active = self.txn_mgr.lock().is_active();
         if is_active {
-            let maybe_table = self
-                .txn_mgr
+            self.txn_mgr
                 .lock()
                 .get_stream_table_source(stream_desc)
-                .map(|table_info| self.get_table_by_info(&table_info));
-            if let Some(t) = maybe_table {
-                return t;
-            }
-            self.get_table(tenant, db_name, source_table_name).await
+                .map(|table_info| self.get_table_by_info(&table_info))
+                .transpose()
         } else {
-            self.get_table(tenant, db_name, source_table_name).await
+            Ok(None)
+        }
+    }
+
+    // Cache stream source table to buffer.
+    fn cache_stream_source_table(&self, stream: TableInfo, source: TableInfo) {
+        let is_active = self.txn_mgr.lock().is_active();
+        if is_active {
+            self.txn_mgr.lock().upsert_stream_table(stream, source);
         }
     }
 }
diff --git a/src/query/catalog/src/table.rs b/src/query/catalog/src/table.rs
@@ -39,7 +39,6 @@ use databend_storages_common_table_meta::meta::TableSnapshot;
 use databend_storages_common_table_meta::table::ChangeType;
 use databend_storages_common_table_meta::table::StreamMode;
 
-use crate::catalog::Catalog;
 use crate::lock::Lock;
 use crate::plan::DataSourceInfo;
 use crate::plan::DataSourcePlan;
@@ -402,14 +401,6 @@ pub trait Table: Sync + Send {
     fn is_read_only(&self) -> bool {
         false
     }
-
-    async fn stream_source_table(&self, _catalog: Arc<dyn Catalog>) -> Result<Arc<dyn Table>> {
-        Err(ErrorCode::Unimplemented(format!(
-            "The 'stream_source_table' operation is not supported for the table '{}'. Table engine: '{}'.",
-            self.name(),
-            self.get_table_info().engine(),
-        )))
-    }
 }
 
 #[async_trait::async_trait]
diff --git a/src/query/service/src/interpreters/common/stream.rs b/src/query/service/src/interpreters/common/stream.rs
@@ -48,7 +48,7 @@ pub async fn build_update_stream_meta_seq(
     for table in tables.into_iter() {
         let stream = StreamTable::try_from_table(table.as_ref())?;
         let stream_info = stream.get_table_info();
-        let source_table = stream.source_table(ctx.get_default_catalog()?).await?;
+        let source_table = stream.source_table(ctx.clone()).await?;
         let inner_fuse = FuseTable::try_from_table(source_table.as_ref())?;
 
         let table_version = inner_fuse.get_table_info().ident.seq;
diff --git a/src/query/service/src/sessions/query_ctx_shared.rs b/src/query/service/src/sessions/query_ctx_shared.rs
@@ -46,6 +46,7 @@ use databend_common_storage::CopyStatus;
 use databend_common_storage::DataOperator;
 use databend_common_storage::MergeStatus;
 use databend_common_storage::StorageMetrics;
+use databend_common_storages_stream::stream_table::StreamTable;
 use databend_common_users::UserApiProvider;
 use parking_lot::Mutex;
 use parking_lot::RwLock;
@@ -309,15 +310,18 @@ impl QueryContextShared {
         let table_meta_key = (catalog.to_string(), database.to_string(), table.to_string());
 
         let already_in_cache = { self.tables_refs.lock().contains_key(&table_meta_key) };
-        match already_in_cache {
-            false => self.get_table_to_cache(catalog, database, table).await,
-            true => Ok(self
+        let res = match already_in_cache {
+            false => self.get_table_to_cache(catalog, database, table).await?,
+            true => self
                 .tables_refs
                 .lock()
                 .get(&table_meta_key)
                 .ok_or_else(|| ErrorCode::Internal("Logical error, it's a bug."))?
-                .clone()),
-        }
+                .clone(),
+        };
+
+        self.cache_stream_source_table(res.clone(), catalog).await?;
+        Ok(res)
     }
 
     #[async_backtrace::framed]
@@ -343,6 +347,47 @@ impl QueryContextShared {
         }
     }
 
+    // Cache the source table of a stream table to ensure can get the same table metadata.
+    #[async_backtrace::framed]
+    async fn cache_stream_source_table(&self, table: Arc<dyn Table>, catalog: &str) -> Result<()> {
+        if table.engine() == "STREAM" {
+            let stream = StreamTable::try_from_table(table.as_ref())?;
+            let table_name = stream.source_table_name();
+            let database = stream.source_table_database();
+            let meta_key = (
+                catalog.to_string(),
+                database.to_string(),
+                table_name.to_string(),
+            );
+            let already_in_cache = { self.tables_refs.lock().contains_key(&meta_key) };
+            if !already_in_cache {
+                let stream_desc = &stream.get_table_info().desc;
+                let tenant = self.get_tenant();
+                let catalog = self
+                    .catalog_manager
+                    .get_catalog(tenant.as_str(), catalog, self.session.session_ctx.txn_mgr())
+                    .await?;
+                let source_table = match catalog.get_stream_source_table(stream_desc)? {
+                    Some(source_table) => source_table,
+                    None => {
+                        let source_table = catalog
+                            .get_table(tenant.as_str(), database, table_name)
+                            .await?;
+                        catalog.cache_stream_source_table(
+                            stream.get_table_info().clone(),
+                            source_table.get_table_info().clone(),
+                        );
+                        source_table
+                    }
+                };
+
+                let mut tables_refs = self.tables_refs.lock();
+                tables_refs.entry(meta_key).or_insert(source_table.clone());
+            }
+        }
+        Ok(())
+    }
+
     pub fn evict_table_from_cache(&self, catalog: &str, database: &str, table: &str) -> Result<()> {
         let table_meta_key = (catalog.to_string(), database.to_string(), table.to_string());
         let mut tables_refs = self.tables_refs.lock();
diff --git a/src/query/storages/common/txn/src/manager.rs b/src/query/storages/common/txn/src/manager.rs
@@ -152,15 +152,18 @@ impl TxnManager {
         self.txn_buffer.update_table_meta(req, table_info);
     }
 
-    pub fn add_stream_table(&mut self, stream: TableInfo, source: TableInfo) {
-        self.txn_buffer
-            .table_desc_to_id
-            .insert(stream.desc.clone(), stream.ident.table_id);
+    pub fn upsert_stream_table(&mut self, stream: TableInfo, source: TableInfo) {
         self.txn_buffer
             .stream_tables
             .insert(stream.ident.table_id, StreamSnapshot { stream, source });
     }
 
+    pub fn upsert_table_desc_to_id(&mut self, table: TableInfo) {
+        self.txn_buffer
+            .table_desc_to_id
+            .insert(table.desc.clone(), table.ident.table_id);
+    }
+
     pub fn get_stream_table_source(&self, stream_desc: &str) -> Option<TableInfo> {
         self.txn_buffer
             .table_desc_to_id
diff --git a/src/query/storages/stream/src/stream_table.rs b/src/query/storages/stream/src/stream_table.rs
@@ -17,7 +17,6 @@ use std::collections::HashSet;
 use std::sync::Arc;
 use std::time::Instant;
 
-use databend_common_catalog::catalog::Catalog;
 use databend_common_catalog::catalog::StorageDescription;
 use databend_common_catalog::plan::block_id_from_location;
 use databend_common_catalog::plan::DataSourcePlan;
@@ -133,11 +132,10 @@ impl StreamTable {
         })
     }
 
-    pub async fn source_table(&self, catalog: Arc<dyn Catalog>) -> Result<Arc<dyn Table>> {
-        let table = catalog
-            .stream_source_table(
-                &self.stream_info.desc,
-                &self.stream_info.tenant,
+    pub async fn source_table(&self, ctx: Arc<dyn TableContext>) -> Result<Arc<dyn Table>> {
+        let table = ctx
+            .get_table(
+                self.stream_info.catalog(),
                 &self.table_database,
                 &self.table_name,
             )
@@ -192,7 +190,7 @@ impl StreamTable {
         push_downs: Option<PushDownInfo>,
     ) -> Result<(PartStatistics, Partitions)> {
         let start = Instant::now();
-        let table = self.source_table(ctx.get_default_catalog()?).await?;
+        let table = self.source_table(ctx.clone()).await?;
         let fuse_table = FuseTable::try_from_table(table.as_ref())?;
 
         let fuse_segment_io =
@@ -288,7 +286,7 @@ impl StreamTable {
 
     #[minitrace::trace]
     pub async fn check_stream_status(&self, ctx: Arc<dyn TableContext>) -> Result<StreamStatus> {
-        let base_table = self.source_table(ctx.get_default_catalog()?).await?;
+        let base_table = self.source_table(ctx).await?;
         let status = if base_table.get_table_info().ident.seq == self.table_version {
             StreamStatus::NoData
         } else {
@@ -343,7 +341,7 @@ impl Table for StreamTable {
                     return Ok(StreamMode::AppendOnly);
                 }
 
-                let table = self.source_table(ctx.get_default_catalog()?).await?;
+                let table = self.source_table(ctx).await?;
                 let fuse_table = FuseTable::try_from_table(table.as_ref())?;
 
                 let latest_snapshot = fuse_table.read_table_snapshot().await?;
@@ -404,7 +402,7 @@ impl Table for StreamTable {
         ctx: Arc<dyn TableContext>,
         change_type: Option<ChangeType>,
     ) -> Result<Option<TableStatistics>> {
-        let table = self.source_table(ctx.get_default_catalog()?).await?;
+        let table = self.source_table(ctx.clone()).await?;
         let fuse_table = FuseTable::try_from_table(table.as_ref())?;
         let change_type = change_type.unwrap_or(ChangeType::Append);
 
@@ -485,13 +483,9 @@ impl Table for StreamTable {
         &self,
         ctx: Arc<dyn TableContext>,
     ) -> Result<Box<dyn ColumnStatisticsProvider>> {
-        let table = self.source_table(ctx.get_default_catalog()?).await?;
+        let table = self.source_table(ctx.clone()).await?;
         table.column_statistics_provider(ctx).await
     }
-
-    async fn stream_source_table(&self, catalog: Arc<dyn Catalog>) -> Result<Arc<dyn Table>> {
-        self.source_table(catalog).await
-    }
 }
 
 fn replace_push_downs(
diff --git a/src/query/storages/system/src/streams_table.rs b/src/query/storages/system/src/streams_table.rs
@@ -194,7 +194,7 @@ impl AsyncSystemTable for StreamsTable {
                         snapshot_location.push(stream_table.snapshot_loc());
 
                         let mut reason = "".to_string();
-                        match stream_table.source_table(ctx.get_default_catalog()?).await {
+                        match stream_table.source_table(ctx.clone()).await {
                             Ok(source) => {
                                 let fuse_table = FuseTable::try_from_table(source.as_ref())?;
                                 if let Some(location) = stream_table.snapshot_loc() {
diff --git a/tests/suites/5_ee/05_stream/05_0001_ee_stream_consume.result b/tests/suites/5_ee/05_stream/05_0001_ee_stream_consume.result
@@ -0,0 +1 @@
+test stream success
diff --git a/tests/suites/5_ee/05_stream/05_0001_ee_stream_consume.sh b/tests/suites/5_ee/05_stream/05_0001_ee_stream_consume.sh
@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+echo "drop database if exists db_stream" | $BENDSQL_CLIENT_CONNECT
+echo "create database db_stream" | $BENDSQL_CLIENT_CONNECT
+echo "create table db_stream.base(a int) change_tracking = true" | $BENDSQL_CLIENT_CONNECT
+echo "create table db_stream.rand like db_stream.base Engine = Random" | $BENDSQL_CLIENT_CONNECT
+echo "create stream db_stream.s on table db_stream.base" | $BENDSQL_CLIENT_CONNECT
+echo "create table db_stream.sink(a int)" | $BENDSQL_CLIENT_CONNECT
+
+# Define function to write data into the base table.
+write_to_base() {
+  for i in {1..20}; do
+    echo "insert into db_stream.base select * from db_stream.rand limit 10" | $BENDSQL_CLIENT_CONNECT
+
+    if (( i % 5 == 0 )); then
+      echo "optimize table db_stream.base compact" | $BENDSQL_CLIENT_CONNECT
+    fi
+  done
+}
+
+# Define function to consume data from the stream into the sink table.
+consume_from_stream() {
+  for i in {1..10}; do
+    echo "insert into db_stream.sink select a from db_stream.s" | $BENDSQL_CLIENT_CONNECT
+  done
+}
+
+# Start the write and consume operations in parallel
+write_to_base &
+write_pid=$!
+consume_from_stream &
+consume_pid=$!
+
+# Wait for the data writing into the base table to complete
+wait $write_pid
+# Wait for the final consume operation to complete
+wait $consume_pid
+
+# Perform a final consume operation from the stream to ensure all data is consumed
+echo "insert into db_stream.sink select a from db_stream.s" | $BENDSQL_CLIENT_CONNECT
+
+# Fetch the counts and sums from both base and sink tables
+base_count=$(echo "SELECT COUNT(*) FROM db_stream.base;" | $BENDSQL_CLIENT_CONNECT)
+sink_count=$(echo "SELECT COUNT(*) FROM db_stream.sink;" | $BENDSQL_CLIENT_CONNECT)
+base_sum=$(echo "SELECT SUM(a) FROM db_stream.base;" | $BENDSQL_CLIENT_CONNECT)
+sink_sum=$(echo "SELECT SUM(a) FROM db_stream.sink;" | $BENDSQL_CLIENT_CONNECT)
+
+# Compare the counts and sums to verify consistency between the base and sink tables
+if [ "$base_count" -eq "$sink_count" ] && [ "$base_sum" -eq "$sink_sum" ]; then
+  echo "test stream success"
+fi
+
+echo "drop stream if exists db_stream.s" | $BENDSQL_CLIENT_CONNECT
+echo "drop table if exists db_stream.base all" | $BENDSQL_CLIENT_CONNECT
+echo "drop table if exists db_stream.rand all" | $BENDSQL_CLIENT_CONNECT
+echo "drop table if exists db_stream.sink all" | $BENDSQL_CLIENT_CONNECT
+echo "drop database if exists db_stream" | $BENDSQL_CLIENT_CONNECT
