From c1792d7be888030bd113c0ca97abe242d99528e7 Mon Sep 17 00:00:00 2001 From: Pamela Fox Date: Fri, 21 Mar 2025 12:39:35 -0700 Subject: [PATCH 1/4] Add cosmos migration files --- scripts/cosmosdb_migration.py | 162 ++++++++++++++++++++++++++ tests/test_cosmosdb_migration.py | 192 +++++++++++++++++++++++++++++++ 2 files changed, 354 insertions(+) create mode 100644 scripts/cosmosdb_migration.py create mode 100644 tests/test_cosmosdb_migration.py diff --git a/scripts/cosmosdb_migration.py b/scripts/cosmosdb_migration.py new file mode 100644 index 0000000000..009231efc2 --- /dev/null +++ b/scripts/cosmosdb_migration.py @@ -0,0 +1,162 @@ +""" +A migration script to migrate data from CosmosDB to a new format. +The old schema: +id: str +entra_oid: str +title: str +timestamp: int +answers: list of 2-item list of str, dict + +The new schema has two item types in the same container: +For session items: +id: str +session_id: str +entra_oid: str +title: str +timestamp: int +type: str (always "session") + +For message_pair items: +id: str +session_id: str +entra_oid: str +type: str (always "message_pair") +version: str (always "cosmosdb-v2") +question: str +response: dict +""" + +import os + +from azure.cosmos.aio import CosmosClient +from azure.identity.aio import AzureDeveloperCliCredential + +from load_azd_env import load_azd_env + + +class CosmosDBMigrator: + """ + Migrator class for CosmosDB data migration. + """ + + def __init__(self, cosmos_account, database_name, credential=None): + """ + Initialize the migrator with CosmosDB account and database. + + Args: + cosmos_account: CosmosDB account name + database_name: Database name + credential: Azure credential, defaults to AzureDeveloperCliCredential + """ + self.cosmos_account = cosmos_account + self.database_name = database_name + self.credential = credential or AzureDeveloperCliCredential() + self.client = None + self.database = None + self.old_container = None + self.new_container = None + + async def connect(self): + """ + Connect to CosmosDB and initialize containers. + """ + self.client = CosmosClient( + url=f"https://{self.cosmos_account}.documents.azure.com:443/", credential=self.credential + ) + self.database = self.client.get_database_client(self.database_name) + self.old_container = self.database.get_container_client("chat-history") + self.new_container = self.database.get_container_client("chat-history-v2") + # check to see that both containers actually exist + try: + await self.old_container.read() + except Exception: + raise ValueError(f"Old container {self.old_container.id} does not exist") + try: + await self.new_container.read() + except Exception: + raise ValueError(f"New container {self.new_container.id} does not exist") + + async def migrate(self): + """ + Migrate data from old schema to new schema. + """ + if not self.client: + await self.connect() + + # Migration: read from old_container, transform, and insert into new_container + query_results = self.old_container.query_items(query="SELECT * FROM c") + + item_migration_count = 0 + async for page in query_results.by_page(): + async for old_item in page: + batch_operations = [] + # Build session item operation + session_item = { + "id": old_item["id"], + "version": "cosmosdb-v2", + "session_id": old_item["id"], + "entra_oid": old_item["entra_oid"], + "title": old_item.get("title"), + "timestamp": old_item.get("timestamp"), + "type": "session", + } + batch_operations.append(("upsert", (session_item,))) + + # Build message_pair operations + answers = old_item.get("answers", []) + for idx, answer in enumerate(answers): + question = answer[0] + response = answer[1] + message_pair = { + "id": f"{old_item['id']}-{idx}", + "version": "cosmosdb-v2", + "session_id": old_item["id"], + "entra_oid": old_item["entra_oid"], + "type": "message_pair", + "question": question, + "response": response, + "order": idx, + "timestamp": None, + } + batch_operations.append(("upsert", (message_pair,))) + + # Execute the batch using partition key [entra_oid, session_id] + await self.new_container.execute_item_batch( + batch_operations=batch_operations, partition_key=[old_item["entra_oid"], old_item["id"]] + ) + item_migration_count += 1 + print(f"Total items migrated: {item_migration_count}") + + async def close(self): + """ + Close the CosmosDB client. + """ + if self.client: + await self.client.close() + + +async def migrate_cosmosdb_data(): + """ + Legacy function for backward compatibility. + Migrate data from CosmosDB to a new format. + """ + # Setup CosmosDB client using configuration similar to cosmosdb.py + USE_CHAT_HISTORY_COSMOS = os.getenv("USE_CHAT_HISTORY_COSMOS", "").lower() == "true" + if not USE_CHAT_HISTORY_COSMOS: + raise ValueError("USE_CHAT_HISTORY_COSMOS must be set to true") + AZURE_COSMOSDB_ACCOUNT = os.environ["AZURE_COSMOSDB_ACCOUNT"] + AZURE_CHAT_HISTORY_DATABASE = os.environ["AZURE_CHAT_HISTORY_DATABASE"] + + migrator = CosmosDBMigrator(AZURE_COSMOSDB_ACCOUNT, AZURE_CHAT_HISTORY_DATABASE) + try: + await migrator.migrate() + finally: + await migrator.close() + + +if __name__ == "__main__": + load_azd_env() + + import asyncio + + asyncio.run(migrate_cosmosdb_data()) diff --git a/tests/test_cosmosdb_migration.py b/tests/test_cosmosdb_migration.py new file mode 100644 index 0000000000..c7d4464a7d --- /dev/null +++ b/tests/test_cosmosdb_migration.py @@ -0,0 +1,192 @@ +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from scripts.cosmosdb_migration import CosmosDBMigrator, migrate_cosmosdb_data + +# Sample old format item +TEST_OLD_ITEM = { + "id": "123", + "entra_oid": "OID_X", + "title": "This is a test message", + "timestamp": 123456789, + "answers": [ + [ + "What does a Product Manager do?", + { + "delta": {"role": "assistant"}, + "session_state": "143c0240-b2ee-4090-8e90-2a1c58124894", + "message": { + "content": "A Product Manager is responsible for product strategy and execution.", + "role": "assistant", + }, + }, + ], + [ + "What about a Software Engineer?", + { + "delta": {"role": "assistant"}, + "session_state": "243c0240-b2ee-4090-8e90-2a1c58124894", + "message": { + "content": "A Software Engineer writes code to create applications.", + "role": "assistant", + }, + }, + ], + ], +} + + +class MockAsyncPageIterator: + """Helper class to mock an async page from CosmosDB""" + + def __init__(self, items): + self.items = items + + def __aiter__(self): + return self + + async def __anext__(self): + if not self.items: + raise StopAsyncIteration + return self.items.pop(0) + + +class MockCosmosDBResultsIterator: + """Helper class to mock a paginated query result from CosmosDB""" + + def __init__(self, data=[]): + self.data = data + self.continuation_token = None + + def by_page(self, continuation_token=None): + """Return a paged iterator""" + self.continuation_token = "next_token" if not continuation_token else continuation_token + "_next" + # Return an async iterator that contains pages + return MockPagesAsyncIterator(self.data) + + +class MockPagesAsyncIterator: + """Helper class to mock an iterator of pages""" + + def __init__(self, data): + self.data = data + self.continuation_token = "next_token" + + def __aiter__(self): + return self + + async def __anext__(self): + if not self.data: + raise StopAsyncIteration + # Return a page, which is an async iterator of items + return MockAsyncPageIterator([self.data.pop(0)]) + + +@pytest.mark.asyncio +async def test_migrate_method(): + """Test the migrate method of CosmosDBMigrator""" + # Create mock objects + mock_container = MagicMock() + mock_database = MagicMock() + mock_client = MagicMock() + + # Set up the query_items mock to return our test item + mock_container.query_items.return_value = MockCosmosDBResultsIterator([TEST_OLD_ITEM]) + + # Set up execute_item_batch as a spy to capture calls + execute_batch_mock = AsyncMock() + mock_container.execute_item_batch = execute_batch_mock + + # Set up the database mock to return our container mocks + mock_database.get_container_client.side_effect = lambda container_name: mock_container + + # Set up the client mock + mock_client.get_database_client.return_value = mock_database + + # Create the migrator with our mocks + migrator = CosmosDBMigrator("dummy_account", "dummy_db") + migrator.client = mock_client + migrator.database = mock_database + migrator.old_container = mock_container + migrator.new_container = mock_container + + # Call the migrate method + await migrator.migrate() + + # Verify query_items was called with the right parameters + mock_container.query_items.assert_called_once_with(query="SELECT * FROM c") + + # Verify execute_item_batch was called + execute_batch_mock.assert_called_once() + + # Extract the arguments from the call + call_args = execute_batch_mock.call_args[1] + batch_operations = call_args["batch_operations"] + partition_key = call_args["partition_key"] + + # Verify the partition key + assert partition_key == ["OID_X", "123"] + + # We should have 3 operations: 1 for session and 2 for message pairs + assert len(batch_operations) == 3 + + # Verify session item + session_operation = batch_operations[0] + assert session_operation[0] == "upsert" + session_item = session_operation[1][0] + assert session_item["id"] == "123" + assert session_item["session_id"] == "123" + assert session_item["entra_oid"] == "OID_X" + assert session_item["title"] == "This is a test message" + assert session_item["timestamp"] == 123456789 + assert session_item["type"] == "session" + assert session_item["version"] == "cosmosdb-v2" + + # Verify first message pair + message1_operation = batch_operations[1] + assert message1_operation[0] == "upsert" + message1_item = message1_operation[1][0] + assert message1_item["id"] == "123-0" + assert message1_item["session_id"] == "123" + assert message1_item["entra_oid"] == "OID_X" + assert message1_item["question"] == "What does a Product Manager do?" + assert message1_item["type"] == "message_pair" + assert message1_item["order"] == 0 + + # Verify second message pair + message2_operation = batch_operations[2] + assert message2_operation[0] == "upsert" + message2_item = message2_operation[1][0] + assert message2_item["id"] == "123-1" + assert message2_item["session_id"] == "123" + assert message2_item["entra_oid"] == "OID_X" + assert message2_item["question"] == "What about a Software Engineer?" + assert message2_item["type"] == "message_pair" + assert message2_item["order"] == 1 + + +@pytest.mark.asyncio +async def test_migrate_cosmosdb_data(): + """Test the main migrate_cosmosdb_data function""" + # Set required environment variables + os.environ["USE_CHAT_HISTORY_COSMOS"] = "true" + os.environ["AZURE_COSMOSDB_ACCOUNT"] = "dummy_account" + os.environ["AZURE_CHAT_HISTORY_DATABASE"] = "dummy_db" + + # Create a mock for the CosmosDBMigrator + with patch("scripts.cosmosdb_migration.CosmosDBMigrator") as mock_migrator_class: + # Set up the mock for the migrator instance + mock_migrator = AsyncMock() + mock_migrator_class.return_value = mock_migrator + + # Call the function + await migrate_cosmosdb_data() + + # Verify the migrator was created with the right parameters + mock_migrator_class.assert_called_once_with("dummy_account", "dummy_db") + + # Verify migrate and close were called + mock_migrator.migrate.assert_called_once() + mock_migrator.close.assert_called_once() From 4ca014db90e4af4ba08fb9a669ff5dbe922213ed Mon Sep 17 00:00:00 2001 From: Pamela Fox Date: Fri, 21 Mar 2025 12:41:54 -0700 Subject: [PATCH 2/4] Migrate to CosmosDB --- scripts/cosmosdb_migration.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/scripts/cosmosdb_migration.py b/scripts/cosmosdb_migration.py index 009231efc2..cd62dd6b43 100644 --- a/scripts/cosmosdb_migration.py +++ b/scripts/cosmosdb_migration.py @@ -15,6 +15,7 @@ title: str timestamp: int type: str (always "session") +version: str (always "cosmosdb-v2") For message_pair items: id: str @@ -66,7 +67,6 @@ async def connect(self): self.database = self.client.get_database_client(self.database_name) self.old_container = self.database.get_container_client("chat-history") self.new_container = self.database.get_container_client("chat-history-v2") - # check to see that both containers actually exist try: await self.old_container.read() except Exception: @@ -83,14 +83,13 @@ async def migrate(self): if not self.client: await self.connect() - # Migration: read from old_container, transform, and insert into new_container query_results = self.old_container.query_items(query="SELECT * FROM c") item_migration_count = 0 async for page in query_results.by_page(): async for old_item in page: batch_operations = [] - # Build session item operation + # Build session item session_item = { "id": old_item["id"], "version": "cosmosdb-v2", @@ -102,7 +101,7 @@ async def migrate(self): } batch_operations.append(("upsert", (session_item,))) - # Build message_pair operations + # Build message_pair answers = old_item.get("answers", []) for idx, answer in enumerate(answers): question = answer[0] @@ -140,7 +139,6 @@ async def migrate_cosmosdb_data(): Legacy function for backward compatibility. Migrate data from CosmosDB to a new format. """ - # Setup CosmosDB client using configuration similar to cosmosdb.py USE_CHAT_HISTORY_COSMOS = os.getenv("USE_CHAT_HISTORY_COSMOS", "").lower() == "true" if not USE_CHAT_HISTORY_COSMOS: raise ValueError("USE_CHAT_HISTORY_COSMOS must be set to true") From d649feaca7e5911f81e7abae5be035cdd11588de Mon Sep 17 00:00:00 2001 From: Pamela Fox Date: Fri, 21 Mar 2025 12:55:32 -0700 Subject: [PATCH 3/4] Make mypy happy --- scripts/cosmosdb_migration.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/cosmosdb_migration.py b/scripts/cosmosdb_migration.py index cd62dd6b43..1eb974ec25 100644 --- a/scripts/cosmosdb_migration.py +++ b/scripts/cosmosdb_migration.py @@ -82,6 +82,8 @@ async def migrate(self): """ if not self.client: await self.connect() + if not self.old_container or not self.new_container: + raise ValueError("Containers do not exist") query_results = self.old_container.query_items(query="SELECT * FROM c") From 1d3540abf971fd1dcd633e463c4a60f97c1793e9 Mon Sep 17 00:00:00 2001 From: Pamela Fox Date: Fri, 21 Mar 2025 14:55:48 -0700 Subject: [PATCH 4/4] Mock env the right way --- tests/test_cosmosdb_migration.py | 42 ++++++++++++++++---------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/test_cosmosdb_migration.py b/tests/test_cosmosdb_migration.py index c7d4464a7d..e67c450061 100644 --- a/tests/test_cosmosdb_migration.py +++ b/tests/test_cosmosdb_migration.py @@ -168,25 +168,25 @@ async def test_migrate_method(): @pytest.mark.asyncio -async def test_migrate_cosmosdb_data(): +async def test_migrate_cosmosdb_data(monkeypatch): """Test the main migrate_cosmosdb_data function""" - # Set required environment variables - os.environ["USE_CHAT_HISTORY_COSMOS"] = "true" - os.environ["AZURE_COSMOSDB_ACCOUNT"] = "dummy_account" - os.environ["AZURE_CHAT_HISTORY_DATABASE"] = "dummy_db" - - # Create a mock for the CosmosDBMigrator - with patch("scripts.cosmosdb_migration.CosmosDBMigrator") as mock_migrator_class: - # Set up the mock for the migrator instance - mock_migrator = AsyncMock() - mock_migrator_class.return_value = mock_migrator - - # Call the function - await migrate_cosmosdb_data() - - # Verify the migrator was created with the right parameters - mock_migrator_class.assert_called_once_with("dummy_account", "dummy_db") - - # Verify migrate and close were called - mock_migrator.migrate.assert_called_once() - mock_migrator.close.assert_called_once() + with patch.dict(os.environ, clear=True): + monkeypatch.setenv("USE_CHAT_HISTORY_COSMOS", "true") + monkeypatch.setenv("AZURE_COSMOSDB_ACCOUNT", "dummy_account") + monkeypatch.setenv("AZURE_CHAT_HISTORY_DATABASE", "dummy_db") + + # Create a mock for the CosmosDBMigrator + with patch("scripts.cosmosdb_migration.CosmosDBMigrator") as mock_migrator_class: + # Set up the mock for the migrator instance + mock_migrator = AsyncMock() + mock_migrator_class.return_value = mock_migrator + + # Call the function + await migrate_cosmosdb_data() + + # Verify the migrator was created with the right parameters + mock_migrator_class.assert_called_once_with("dummy_account", "dummy_db") + + # Verify migrate and close were called + mock_migrator.migrate.assert_called_once() + mock_migrator.close.assert_called_once()