db_setup.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os.path
import json
import sys
import time
import argparse
import uuid
import pymongo
import datetime

args = {}

client = None
gwentDB = None

SOURCE_FOLDER = "./"
FILES = {"rarities": SOURCE_FOLDER + "rarities.jsonl", "groups": SOURCE_FOLDER + "groups.jsonl",
 "factions": SOURCE_FOLDER + "factions.jsonl", "categories": SOURCE_FOLDER + "categories.jsonl"}

CARDS_FILE = SOURCE_FOLDER + "cards.jsonl"


# Set the command line parameters.
def set_parser():
 parser = argparse.ArgumentParser(description='This script allows you to perform the initial seed '
 'for the gwentapi db.')
 parser.add_argument('-f', '--force', help='Disable safeguards and force the database to be dropped'
 'before the insertion of the data.',
 action='store_true', required=False)
 parser.add_argument('-u', '--username', help='Username used for authentication')
 parser.add_argument('-p', '--password', help='Indicate that a password is needed for the authenticated user.'
 'Will start an interactive prompt.', action='store_true')
 parser.add_argument('--host', help='The host to connect to.')
 parser.add_argument('--port', help='Port to connect to.')
 parser.add_argument('--authenticationDatabase', help='Database to authenticate to. Where your user account '
 'was created.')

 global args
 args = parser.parse_args()


def is_insert_safe():
 collections = gwentDB.collection_names(include_system_collections=False)
 for collection, path in FILES.items():
 if collection in collections:
 return False
 if "cards" in collections:
 return False
 if "variations" in collections:
 return False
 return True


def insertUUID(item):
 id = uuid.uuid4().bytes
 item["uuid"] = id


def testIfFileExists():
 for key, path in FILES.items():
 if not os.path.isfile(path):
 print("Error: " + path + " doesn't exists.")
 return False
 return True


def drop_collections(collection):
 print("Dropping the '" + collection + "' collection.")
 gwentDB[collection].drop()


def mongodb_conn():
 try:
 if args.host and args.port:
 return pymongo.MongoClient(args.host, args.port)
 elif args.host:
 return pymongo.MongoClient(args.host)
 else:
 return pymongo.MongoClient()
 except pymongo.errors.ConnectionFailure as e:
 print("Could not connect to the server: %s" % e)


# Untested on other platforms other than Linux.
def getpass():
 password = ""
 if sys.platform != "win32":
 os.system("stty -echo")
 password = input("Enter Password: ")
 os.system("stty echo")
 print("")
 else:
 password = input("Enter Password: ")
 return password


def mongodb_auth():
 try:
 if args.password and args.username and args.authenticationDatabase:
 password = getpass()
 client.gwentapi.authenticate(args.username, password, source=args.authenticationDatabase)
 elif args.password and args.username:
 password = getpass()
 client.gwentapi.authenticate(args.username, password)
 return True
 except pymongo.errors.OperationFailure as e:
 print("Error: %s" % e)
 return False


def main():
 global gwentDB
 global client

 client = mongodb_conn()

 gwentDB = client.gwentapi
 if not mongodb_auth():
 return

 if not testIfFileExists():
 return
 if not args.force and not is_insert_safe():
 print("The database already contain some collections.")
 print("The operation was aborted.")
 return

 for collection, path in FILES.items():
 print("Working on: " + collection)
 drop_collections(collection)
 with open(path, encoding="utf-8", newline="\n") as f:
 for line in f:
 data = json.loads(line)
 insertUUID(data)
 data['last_modified'] = datetime.datetime.utcnow()
 gwentDB[collection].insert_one(data)
 gwentDB[collection].create_index([('name', pymongo.ASCENDING)], unique=True)
 gwentDB[collection].create_index([('uuid', pymongo.ASCENDING)], unique=True)
 # Processing cards
 print("Working on: cards")
 drop_collections("cards")
 drop_collections("variations")
 with open(CARDS_FILE, encoding="utf-8", newline="\n") as f:
 for line in f:
 data = json.loads(line)
 for variation in data["variations"]:
 variation["art"]["fullsizeImage"] = data["key"] + "_full.png"
 variation["art"]["thumbnailImage"] = data["key"] + "_thumbnail.png"
 data.pop("key", None)
 # Manual reference for faction
 if "faction" in data:
 faction_id = gwentDB.factions.find_one({"name": data["faction"]})["_id"]
 data["faction_id"] = faction_id
 # Renaming type to group until the jsonl format change.
 # Manual reference for group
 if "type" in data:
 data["group"] = data.pop("type")
 group_id = gwentDB.groups.find_one({"name": data["group"]})["_id"]
 data["group_id"] = group_id
 if "categories" in data:
 categories_id = list()
 for category in data["categories"]:
 category_id = gwentDB.categories.find_one({"name": category})["_id"]
 categories_id.append(category_id)
 data["categories_id"] = categories_id

 insertUUID(data)
 data['last_modified'] = datetime.datetime.utcnow()
 variations = data.pop("variations")

 card_id = gwentDB.cards.insert_one(data).inserted_id
 for variation in variations:
 if "rarity" in variation:
 rarity_id = gwentDB.rarities.find_one({"name": variation["rarity"]})["_id"]
 variation["rarity_id"] = rarity_id
 variation["card_id"] = card_id
 insertUUID(variation)
 variation['last_modified'] = datetime.datetime.utcnow()
 gwentDB.variations.insert_many(variations)
 gwentDB.cards.create_index([('name', pymongo.ASCENDING)], unique=True)
 gwentDB.cards.create_index([('uuid', pymongo.ASCENDING)], unique=True)
 gwentDB.variations.create_index([('card_id', pymongo.ASCENDING)])
 gwentDB.variations.create_index([('uuid', pymongo.ASCENDING)])

if __name__ == '__main__':
 set_parser()
 print("Starting")
 start = time.time()
 main()
 print("Elapsed Time: %s" % (time.time() - start))