DX-114411: MCP Search - Adding by default removeCatalogName to true for Catalog API call (#84)

* DX-114411: MCP Search - Adding by default removeCatalogName to true for Catalog API call

* updated mcp with suggestions

* cleanup

* fix

Author: mthirani2021

src/dremioai/api/dremio/catalog.py

@@ -110,16 +110,21 @@ class LineageResponse(BaseModel):
     children: List[LineageChildren]
 
 
-async def get_lineage(dataset_id_or_path: str) -> Dict[str, Any]:
+async def get_lineage(
+    dataset_id_or_path: str,
+    remove_catalog_name: Optional[bool] = True
+) -> Dict[str, Any]:
     client = AsyncHttpClient()
     if "." in dataset_id_or_path:
         response = await get_schema(dataset_id_or_path, by_id=False)
         dataset_id_or_path = response["id"]
 
     project_id = settings.instance().dremio.project_id
     endpoint = f"/v0/projects/{project_id}/catalog" if project_id else "/api/v3/catalog"
+    params = {"removeCatalogName": str(remove_catalog_name).lower()}
     result: LineageResponse = await client.get(
         f"{endpoint}/{dataset_id_or_path}/graph",
+        params=params,
         deser=LineageResponse,
     )
     return result.model_dump()
@@ -130,10 +135,12 @@ async def get_schema(
     by_id: Optional[bool] = False,
     include_tags: Optional[bool] = False,
     flatten: Optional[bool] = False,
+    remove_catalog_name: Optional[bool] = True
 ) -> Dict[str, Any]:
     client = AsyncHttpClient()
     project_id = settings.instance().dremio.project_id
     endpoint = f"/v0/projects/{project_id}/catalog" if project_id else "/api/v3/catalog"
+    params = {"removeCatalogName": str(remove_catalog_name).lower()}
     if by_id:
         endpoint += "/" + dataset_path_or_id
     else:
@@ -142,7 +149,7 @@ async def get_schema(
                 reader(StringIO(dataset_path_or_id), delimiter=".", dialect=excel)
             )[0]
         endpoint += f'/by-path/{"/".join(dataset_path_or_id)}'
-    schema = await client.get(endpoint)
+    schema = await client.get(endpoint, params=params)
 
     if include_tags:
 

src/dremioai/api/dremio/search.py

@@ -214,7 +214,8 @@ class EnterpriseSearchResultsWrapper(BaseModel):
 
 
 async def get_search_results(
-    search: str | Search, use_df: bool = False
+    search: str | Search, use_df: bool = False,
+    remove_catalog_name: Optional[bool] = True
 ) -> EnterpriseSearchResultsWrapper | pd.DataFrame:
     if isinstance(search, str):
         search = Search(query=search)
@@ -225,11 +226,15 @@ async def get_search_results(
         if settings.instance().dremio.project_id
         else "/api/v3/search"
     )
+
+    params = {"removeCatalogName": str(remove_catalog_name).lower()}
+
     result = []
     response = await client.post(
         endpoint,
         body=search.model_dump(exclude_none=True),
         deser=EnterpriseSearchResults,
+        params=params,
     )
     while response.results and response.error is None and response.more is None:
         result.extend(response.results)
@@ -240,6 +245,7 @@ async def get_search_results(
             endpoint,
             body=search.model_dump(exclude_none=True),
             deser=EnterpriseSearchResults,
+            params=params,
         )
 
     if use_df:

src/dremioai/api/transport.py

@@ -195,11 +195,12 @@ async def post(
         deser: Optional[DeserializationStrategy] = None,
         file: Optional[TextIO] = None,
         top_level_list: bool = False,
+        params: Dict[AnyStr, Any] = None,
     ):
         async with ClientSession(middlewares=(retry_middleware,)) as session:
-            self.log_request("POST", endpoint)
+            self.log_request("POST", endpoint, params)
             async with session.post(
-                f"{self.uri}{endpoint}", headers=self.headers, json=body, ssl=False
+                f"{self.uri}{endpoint}", params=params, headers=self.headers, json=body, ssl=False
             ) as response:
                 return await self.handle_response(
                     response, deser, file, top_level_list=top_level_list