add migration to fix notifications from pre-1.0.0 upgrade
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
+import argparse
+from uuid import UUID
import json
from typing import Callable, Dict, List
from azure.cosmosdb.table.tablebatch import TableBatch
from azure.cosmosdb.table.tableservice import TableService
+from azure.mgmt.storage import StorageManagementClient
+from azure.common.client_factory import get_client_from_cli_profile
def migrate_task_os(table_service: TableService) -> None:
@@ -38,8 +42,38 @@ def migrate_task_os(table_service: TableService) -> None:
print("migrated %s rows" % count)
+def migrate_notification_keys(table_service: TableService) -> None:
+ table_name = "Notification"
+ notifications = table_service.query_entities(
+ table_name, select="PartitionKey,RowKey,config"
+ )
+ partitionKey = None
+
+ count = 0
+ for entry in notifications:
+ try:
+ UUID(entry.PartitionKey)
+ continue
+ except ValueError:
+ pass
+
+ table_service.insert_or_replace_entity(
+ table_name,
+ {
+ "PartitionKey": entry.RowKey,
+ "RowKey": entry.PartitionKey,
+ "config": entry.config,
+ },
+ )
+ table_service.delete_entity(table_name, entry.PartitionKey, entry.RowKey)
+ count += 1
+
+ print("migrated %s rows" % count)
+
+
migrations: Dict[str, Callable[[TableService], None]] = {
- "migrate_task_os": migrate_task_os
+ "migrate_task_os": migrate_task_os,
+ "migrate_notification_keys": migrate_notification_keys,
}
@@ -48,3 +82,26 @@ def migrate(table_service: TableService, migration_names: List[str]) -> None:
print("applying migration '%s'" % name)
migrations[name](table_service)
print("migration '%s' applied" % name)
+
+
+def main():
+ formatter = argparse.ArgumentDefaultsHelpFormatter
+ parser = argparse.ArgumentParser(formatter_class=formatter)
+ parser.add_argument("resource_group")
+ parser.add_argument("storage_account")
+ parser.add_argument("migration", choices=migrations.keys(), nargs="+")
+ args = parser.parse_args()
+
+ client = get_client_from_cli_profile(StorageManagementClient)
+ storage_keys = client.storage_accounts.list_keys(
+ args.resource_group, args.storage_account
+ )
+ table_service = TableService(
+ account_name=args.storage_account, account_key=storage_keys.keys[0].value
+ )
+ print(args.migration)
+ migrate(table_service, args.migration)
+
+
+if __name__ == "__main__":
+ main()
