Repository: risingwavelabs/risingwave
File: src/meta/src/controller/fragment.rs (Rust)

Code before:
        let inner = self.inner.read().await;
        let txn = inner.db.begin().await?;

        let mut query = StreamingJob::find()
            .select_only()
            .column(streaming_job::Column::JobId);

        if let Some(database_id) = database_id {
            query = query
                .join(JoinType::InnerJoin, streaming_job::Relation::Object.def())
                .filter(object::Column::DatabaseId.eq(database_id));
        }

        let jobs: Vec<JobId> = query.into_tuple().all(&txn).await?;

        let jobs: HashSet<JobId> = jobs.into_iter().collect();

        if jobs.is_empty() {
            return Ok(LoadedFragmentContext::default());
        }

        load_fragment_context_for_jobs(&txn, jobs).await
    }

    #[await_tree::instrument]
    pub async fn fill_snapshot_backfill_epoch(
        &self,
        fragment_ids: impl Iterator<Item = FragmentId>,
        snapshot_backfill_info: Option<&SnapshotBackfillInfo>,
        cross_db_snapshot_backfill_info: &SnapshotBackfillInfo,
    ) -> MetaResult<()> {
        let inner = self.inner.write().await;
        let txn = inner.db.begin().await?;
        for fragment_id in fragment_ids {
            let fragment = FragmentModel::find_by_id(fragment_id)
                .one(&txn)
                .await?
                .context(format!("fragment {} not found", fragment_id))?;
            let mut node = fragment.stream_node.to_protobuf();
            if crate::stream::fill_snapshot_backfill_epoch(
                &mut node,
                snapshot_backfill_info,
                cross_db_snapshot_backfill_info,
            )? {
                let node = StreamNode::from(&node);
                FragmentModel::update(fragment::ActiveModel {
                    fragment_id: Set(fragment_id),
                    stream_node: Set(node),
                    ..Default::default()
                })
                .exec(&txn)

Review comment (suggestion):
The new public method `load_fragment_context_in_txn` lacks rustdoc documentation. This is a public API that accepts a transaction handle and should include documentation explaining:
- The purpose of the transactional variant
- Parameters (the generic `txn` parameter type requirements and `database_id` parameter)
- Return value and what LoadedFragmentContext contains
- When to use this transactional variant vs the regular `load_fragment_context` method

Code after:
        )?;

        tracing::trace!(?fragments, "reload all actors");

        Ok(fragments)
    }

    /// Async load stage: collects all metadata required for rendering actor assignments.
    pub async fn load_fragment_context(
        &self,
        database_id: Option<DatabaseId>,
    ) -> MetaResult<LoadedFragmentContext> {
        let inner = self.inner.read().await;
        let txn = inner.db.begin().await?;

        self.load_fragment_context_in_txn(&txn, database_id).await
    }

    pub async fn load_fragment_context_in_txn<C>(
        &self,
        txn: &C,
        database_id: Option<DatabaseId>,
    ) -> MetaResult<LoadedFragmentContext>
    where
        C: ConnectionTrait,
    {
        let mut query = StreamingJob::find()
            .select_only()
            .column(streaming_job::Column::JobId);

        if let Some(database_id) = database_id {
            query = query
                .join(JoinType::InnerJoin, streaming_job::Relation::Object.def())
                .filter(object::Column::DatabaseId.eq(database_id));
        }

        let jobs: Vec<JobId> = query.into_tuple().all(txn).await?;

        let jobs: HashSet<JobId> = jobs.into_iter().collect();

        if jobs.is_empty() {
            return Ok(LoadedFragmentContext::default());
        }

        load_fragment_context_for_jobs(txn, jobs).await
    }

    #[await_tree::instrument]
    pub async fn fill_snapshot_backfill_epoch(
        &self,
        fragment_ids: impl Iterator<Item = FragmentId>,