88import traceback
99from ast import literal_eval
1010from azure .cosmos import CosmosClient
11+ from azure .identity import AzurePowerShellCredential , ChainedTokenCredential , AzureCliCredential
1112from azure .storage .blob import BlobServiceClient
1213
1314logging .getLogger ().setLevel (logging .INFO )
2526BACKUP_CONTAINER = "backups"
2627BLOB_NAME_PATTERN = "cosmos/{0}/{1}"
2728
29+ # Create a AzurePowerShellCredential()
30+ credential_chain = ChainedTokenCredential (AzureCliCredential (), AzurePowerShellCredential ())
2831
32+ def restore_data_from_backup (backup_storage_url , dest_url , db_name ):
2933
30- def restore_data_from_backup (connection_string , dest_url , dest_key , db_name ):
31-
32- dest_db_client = get_db_client (dest_url , dest_key , db_name )
34+ dest_db_client = get_db_client (dest_url , db_name )
3335
34- blob_service_client = BlobServiceClient . from_connection_string ( connection_string )
36+ blob_service_client = BlobServiceClient ( backup_storage_url , credential = credential_chain )
3537 container_client = blob_service_client .get_container_client (BACKUP_CONTAINER )
3638 for cosmos_container_name in COSMOS_CONTAINERS :
3739 # Load source records from backup file
@@ -68,10 +70,10 @@ def get_backup_contents(container_client, blob_name):
6870
6971
7072# Create cosmosdb clients
71- def get_db_client (dest_url , dest_key , db_name ):
73+ def get_db_client (dest_url , db_name ):
7274
7375 # Create cosmosdb client for destination db
74- dest_cosmos_client = CosmosClient (dest_url , credential = dest_key )
76+ dest_cosmos_client = CosmosClient (dest_url , credential = credential_chain )
7577 if not dest_cosmos_client :
7678 logging .error ("Failed to create cosmos client for destination db" )
7779 exit (1 )
@@ -83,7 +85,7 @@ def get_db_client(dest_url, dest_key, db_name):
8385 dest_db_client = dest_cosmos_client .get_database_client (db_name )
8486 logging .info ("Created database clients" )
8587 except :
86- logging .error ("Failed to create databae client using CosmosClient" )
88+ logging .error ("Failed to create database client using CosmosClient" )
8789 traceback .print_exc ()
8890 exit (1 )
8991 return dest_db_client
@@ -122,20 +124,15 @@ def fetch_records(container_client, container_name):
122124 )
123125
124126 parser .add_argument (
125- "--backup-connection-string " ,
127+ "--backup-storage-url " ,
126128 required = True ,
127- help = ("Connection string to backup storage account" ),
129+ help = ("URL to backup storage account" ),
128130 )
129131 parser .add_argument (
130132 "--dest-url" ,
131133 required = True ,
132134 help = ("URL to destination cosmosdb" ),
133135 )
134- parser .add_argument (
135- "--dest-key" ,
136- required = True ,
137- help = ("Destination cosmosdb account key" ),
138- )
139136 parser .add_argument (
140137 "--db-name" ,
141138 required = True ,
@@ -145,4 +142,4 @@ def fetch_records(container_client, container_name):
145142 args = parser .parse_args ()
146143
147144 logging .info ("Syncing database.." )
148- restore_data_from_backup (args .backup_connection_string , args .dest_url , args . dest_key , args .db_name )
145+ restore_data_from_backup (args .backup_storage_url , args .dest_url , args .db_name )
0 commit comments