80 строки
2.7 KiB
Python
80 строки
2.7 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import uuid
|
|
import argparse
|
|
from azure.core.exceptions import AzureError
|
|
from azure.cosmos import CosmosClient, PartitionKey
|
|
import azure.cosmos.exceptions as exceptions
|
|
|
|
# Get input arguments
|
|
parser = argparse.ArgumentParser(description='Translate a JSON file')
|
|
parser.add_argument('--cosmosdb-url', dest='cosmosdb_url', action='store',
|
|
help='You need to supply the URL to access your Cosmos DB')
|
|
parser.add_argument('--cosmosdb-key', dest='cosmosdb_key', action='store',
|
|
help='You need to supply a read/write key to access your CosmosDB')
|
|
parser.add_argument('--input-file', dest='input_file', action='store',
|
|
help='You need to supply file name where the JSON checklist is located')
|
|
parser.add_argument('--discriminator', dest='discriminator', action='store',
|
|
help='Optionally, you can add the a Discriminator field to each document')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Variables
|
|
db_name = "checklist"
|
|
container_name = "checklist"
|
|
|
|
# Get JSON
|
|
try:
|
|
with open(args.input_file) as f:
|
|
checklist = json.load(f)
|
|
except Exception as e:
|
|
print("ERROR: Error when processing JSON file", args.file_name_in, "-", str(e))
|
|
sys.exit(1)
|
|
|
|
# Create client
|
|
try:
|
|
client = CosmosClient(url=args.cosmosdb_url, credential=args.cosmosdb_key)
|
|
except Exception as e:
|
|
print("ERROR: Error when connecting to Cosmos DB", args.cosmosdb_url, "-", str(e))
|
|
sys.exit(1)
|
|
|
|
# Delete database
|
|
try:
|
|
database = client.delete_database(db_name)
|
|
print(f"DEBUG: Database {db_name} deleted")
|
|
except exceptions.CosmosResourceNotFoundError:
|
|
print('INFO: A database with id {db_name} does not exist')
|
|
|
|
# Create database
|
|
try:
|
|
database = client.create_database(id=db_name)
|
|
print(f"DEBUG: Database created: {database.id}")
|
|
except exceptions.CosmosResourceExistsError:
|
|
print("ERROR: Database", db_name, "already exists.")
|
|
|
|
# Create container
|
|
try:
|
|
partition_key_path = PartitionKey(path="/guid", kind="Hash")
|
|
container = database.create_container(
|
|
id=container_name,
|
|
partition_key=partition_key_path
|
|
# offer_throughput=400,
|
|
)
|
|
print(f"DEBUG: Container created: {container.id}")
|
|
except exceptions.CosmosResourceExistsError:
|
|
print("ERROR: Container", container_name, "already exists.")
|
|
|
|
# Upload checklist items
|
|
print ("DEBUG: Adding items, this can take a few minutes...")
|
|
item_counter = 0
|
|
for item in checklist['items']:
|
|
# print("DEBUG: uploading item: {0}".format(str(item)))
|
|
item['id'] = item['guid']
|
|
if args.discriminator:
|
|
item['Discriminator'] = args.discriminator
|
|
container.create_item(body=item)
|
|
item_counter += 1
|
|
|
|
# Finish
|
|
print("{0} items were uploaded to {1}".format(item_counter, args.cosmosdb_url)) |