From 3436793a374b8c6c19963a85257e4d4646e1f487 Mon Sep 17 00:00:00 2001
From: Liubov Dmitrieva <liubovd@meta.com>
Date: Fri, 1 Mar 2024 06:55:56 -0800
Subject: [PATCH] complete annotation for bonsai <=> hg mapping queries with
 e2e

Summary:
complete annotation for bonsai <=> hg mapping queries with e2e
telemetry

we have already covered queries to translate from bonsai to hg, here we are
covering the opposite queries

Reviewed By: mzr

Differential Revision: D54412246

fbshipit-source-id: 3de3bccce9c203facf131fdbe0dedc8a3731c893
---
 eden/mononoke/bonsai_hg_mapping/src/lib.rs | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/eden/mononoke/bonsai_hg_mapping/src/lib.rs b/eden/mononoke/bonsai_hg_mapping/src/lib.rs
index 354eb4fd29..b17bdac0bf 100644
--- a/eden/mononoke/bonsai_hg_mapping/src/lib.rs
+++ b/eden/mononoke/bonsai_hg_mapping/src/lib.rs
@@ -502,10 +502,20 @@ async fn select_mapping(
                     let conn = connection.conn.clone();
                     move |hg_cs_ids| async move {
                         let hg_cs_ids = hg_cs_ids.into_iter().collect::<Vec<_>>();
-                        Ok(SelectMappingByHg::query(&conn, &repo_id, &hg_cs_ids[..])
-                            .await?
-                            .into_iter()
-                            .collect())
+                        Ok(match cri {
+                            Some(cri) => {
+                                SelectMappingByHg::traced_query(
+                                    &conn,
+                                    &cri,
+                                    &repo_id,
+                                    &hg_cs_ids[..],
+                                )
+                                .await
+                            }
+                            None => SelectMappingByHg::query(&conn, &repo_id, &hg_cs_ids[..]).await,
+                        }?
+                        .into_iter()
+                        .collect())
                     }
                 })
                 .await?;