From 52a870ae4eaa5ade40aa181aac34277b25f3e7ec Mon Sep 17 00:00:00 2001 From: Dhruvkumar-Microsoft Date: Fri, 11 Sep 2026 15:48:27 +0530 Subject: [PATCH] Enhance CosmosDB queries to scope results by user ID for improved data security --- src/backend/common/database/cosmosdb.py | 25 +++++++++++-------- .../backend/common/database/test_cosmosdb.py | 11 +++++--- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/src/backend/common/database/cosmosdb.py b/src/backend/common/database/cosmosdb.py index 4eb7e6fe3..1a575effe 100644 --- a/src/backend/common/database/cosmosdb.py +++ b/src/backend/common/database/cosmosdb.py @@ -187,8 +187,8 @@ async def update_plan(self, plan: Plan) -> None: await self.update_item(plan) async def get_plan_by_plan_id(self, plan_id: str) -> Optional[Plan]: - """Retrieve a plan by plan_id.""" - query = "SELECT * FROM c WHERE c.id=@plan_id AND c.data_type=@data_type" + """Retrieve a plan by plan_id, scoped to the current user.""" + query = "SELECT * FROM c WHERE c.id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id" parameters = [ {"name": "@plan_id", "value": plan_id}, {"name": "@data_type", "value": DataType.plan}, @@ -243,11 +243,12 @@ async def update_step(self, step: Step) -> None: await self.update_item(step) async def get_steps_by_plan(self, plan_id: str) -> List[Step]: - """Retrieve all steps for a plan.""" - query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type ORDER BY c.timestamp" + """Retrieve all steps for a plan, scoped to the current user.""" + query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id ORDER BY c.timestamp" parameters = [ {"name": "@plan_id", "value": plan_id}, {"name": "@data_type", "value": DataType.step}, + {"name": "@user_id", "value": self.user_id}, ] return await self.query_items(query, parameters, Step) @@ -441,11 +442,13 @@ async def update_current_team(self, current_team: UserCurrentTeam) -> None: await self.update_item(current_team) async def delete_plan_by_plan_id(self, plan_id: str) -> bool: - """Delete a plan by its ID.""" - query = "SELECT c.id, c.session_id FROM c WHERE c.id=@plan_id " + """Delete a plan by its ID, scoped to the current user.""" + query = "SELECT c.id, c.session_id FROM c WHERE c.id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id" params = [ {"name": "@plan_id", "value": plan_id}, + {"name": "@data_type", "value": DataType.plan}, + {"name": "@user_id", "value": self.user_id}, ] items = self.container.query_items(query=query, parameters=params) self.logger.debug("delete_plan_by_plan_id: querying items for plan_id=%s", plan_id) @@ -471,11 +474,12 @@ async def update_mplan(self, mplan: MPlan) -> None: await self.update_item(mplan) async def get_mplan(self, plan_id: str) -> Optional[MPlan]: - """Retrieve a mplan configuration by mplan_id.""" - query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type" + """Retrieve a mplan configuration by mplan_id, scoped to the current user.""" + query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id" parameters = [ {"name": "@plan_id", "value": plan_id}, {"name": "@data_type", "value": DataType.m_plan}, + {"name": "@user_id", "value": self.user_id}, ] results = await self.query_items(query, parameters, MPlan) return results[0] if results else None @@ -489,11 +493,12 @@ async def update_agent_message(self, message: AgentMessageData) -> None: await self.update_item(message) async def get_agent_messages(self, plan_id: str) -> List[AgentMessageData]: - """Retrieve an agent message by message_id.""" - query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type ORDER BY c._ts ASC" + """Retrieve agent messages for a plan, scoped to the current user.""" + query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id ORDER BY c._ts ASC" parameters = [ {"name": "@plan_id", "value": plan_id}, {"name": "@data_type", "value": DataType.m_plan_message}, + {"name": "@user_id", "value": self.user_id}, ] return await self.query_items(query, parameters, AgentMessageData) diff --git a/src/tests/backend/common/database/test_cosmosdb.py b/src/tests/backend/common/database/test_cosmosdb.py index 8aa901fef..24b718af0 100644 --- a/src/tests/backend/common/database/test_cosmosdb.py +++ b/src/tests/backend/common/database/test_cosmosdb.py @@ -481,7 +481,7 @@ async def test_get_plan_by_plan_id_found(self, client): result = await client.get_plan_by_plan_id("test_plan_id") assert result == mock_plan - expected_query = "SELECT * FROM c WHERE c.id=@plan_id AND c.data_type=@data_type" + expected_query = "SELECT * FROM c WHERE c.id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id" expected_params = [ {"name": "@plan_id", "value": "test_plan_id"}, {"name": "@data_type", "value": DataType.plan}, @@ -608,10 +608,11 @@ async def test_get_steps_by_plan(self, client): result = await client.get_steps_by_plan("test_plan_id") assert result == mock_steps - expected_query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type ORDER BY c.timestamp" + expected_query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id ORDER BY c.timestamp" expected_params = [ {"name": "@plan_id", "value": "test_plan_id"}, {"name": "@data_type", "value": DataType.step}, + {"name": "@user_id", "value": "test_user"}, ] client.query_items.assert_called_once_with(expected_query, expected_params, Step) @@ -963,10 +964,11 @@ async def test_get_agent_messages(self, client): result = await client.get_agent_messages("test_plan_id") assert result == mock_messages - expected_query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type ORDER BY c._ts ASC" + expected_query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id ORDER BY c._ts ASC" expected_params = [ {"name": "@plan_id", "value": "test_plan_id"}, {"name": "@data_type", "value": DataType.m_plan_message}, + {"name": "@user_id", "value": "test_user"}, ] client.query_items.assert_called_once_with(expected_query, expected_params, AgentMessageData) @@ -1038,10 +1040,11 @@ async def test_get_mplan(self, client): result = await client.get_mplan("test_plan_id") assert result == mock_mplan - expected_query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type" + expected_query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type AND c.user_id=@user_id" expected_params = [ {"name": "@plan_id", "value": "test_plan_id"}, {"name": "@data_type", "value": DataType.m_plan}, + {"name": "@user_id", "value": "test_user"}, ] client.query_items.assert_called_once_with(expected_query, expected_params, MPlan)