UploadSales.py

Go to the documentation of this file.

import os
import ssl
import boto3
import awswrangler as wr
import multiprocessing as mp
from multiprocessing.pool import Pool
from itertools import repeat
from pymongo import MongoClient, DeleteOne
from pymongo.errors import BulkWriteError
from MongoConnection import mongo_connection
from dotenv import load_dotenv, dotenv_values
load_dotenv()
 
 
collection_name = os.environ.get("COLLECTION")
mongo_conn = mongo_connection()
 
 
collection = mongo_conn.mongo_conn_sales()
 
catalogue_collection = mongo_conn.mongo_conn_catalogue()
# releases = list(catalogue_collection.find({}, {"_id":0, "upc":1,"title":1,"label":1,"artist":1,"tracks.isrc":1, "tracks.title":1, "tracks.artist":1, "catalogId": 1}))
 
def upload_data_mongo(df):
    """ Uploads batch lines to mongodb
 
    Args:
        df (dataframe): Dataframe with sales matched by catalogue
    Returns: Nothing
    """
    lines = df.apply(lambda x: x.to_dict(), axis=1).to_list()
    if len(lines) > 0:
        collection.insert_many(lines)
    
def find_upc(upcs, df):
    """ Makes catalogue matching using upc as key
 
    Args:
        upcs (list): A list which contains all unique UPCs in the current batch
        df (dataframe): Dataframe with sales matched by catalogue
    Returns: Nothing
    """
    for upc in upcs:
        # upc_found = [element for element in releases if element['upc'] == upc]
        # if len(upc_found) > 0:
        upc_found = catalogue_collection.find_one({"upc":{"$regex":upc}},{"_id":0, "upc":1,"title":1,"label":1,"artist":1, "catalogId": 1})
        ind = df[df["release_id"] == upc].index
        if upc_found:
            df.loc[ind, "release_title"] = upc_found["title"]
            df.loc[ind, "label_id"] = upc_found["label"]
            df.loc[ind, "catalogue_id"] = upc_found["catalogId"]
            df.loc[ind, "clean"] = 1
            df.loc[ind, "release_id"] = upc_found["upc"]
            upload_data_mongo(df.loc[ind,:])
        else:
            upload_data_mongo(df.loc[ind,:])
    # return df
 
def find_irsc(isrc, df):
    """ Makes catalogue matching using isrc as key
 
    Args:
        isrc (list): A list which contains all unique ISRCs in the current batch
        df (dataframe): Dataframe with sales matched by catalogue
    Returns: Nothing
    """
    if len(isrc) < 12:
        aux = df[df["isrc_id"] == isrc]["release_id"].unique()
        find_upc(aux, df)
    else:
        upc_found = catalogue_collection.find_one({"tracks.isrc":isrc},{"_id":0, "upc":1,"title":1,"label":1,"artist":1,"tracks.isrc":1, "tracks.title":1, "tracks.artist":1, "catalogId": 1})
        ind = df[df["isrc_id"] == isrc].index
        if upc_found:
            isrc_found = [element for element in upc_found["tracks"] if element["isrc"] == isrc]
            df.loc[ind, "release_id"] = upc_found["upc"]
            df.loc[ind, "release_title"] = upc_found["title"]
            df.loc[ind, "label_id"] = upc_found["label"]
            df.loc[ind, "track_title"] = isrc_found[0]["title"]
            df.loc[ind, "artists"] = isrc_found[0]["artist"]
            df.loc[ind, "catalogue_id"] = upc_found["catalogId"]
            df.loc[ind, "clean"] = 1
        upload_data_mongo(df.loc[ind,:])
    # return df
 
def delete_files_mongoDB(files):
    """ Deletes all files in the current batch from database (in case of they exist) to avoid overwriting
 
    Args:
        files (list): A list which contains all filenames in the current batch splitted by formats
    Returns: 
        files_delete (list): A list with filenames in the current batch
    """
    files_delete = [i["file"] for i in files]
    r = collection.delete_many({"file": {"$in": files_delete}})
    if r.deleted_count >= 1:
        response = f"{r.deleted_count} sales lines deleted"
    else:
        response = "No files deleted"
    print(response)
    return files_delete
 
def delete_duplicates():
    """ Deletes all duplicated sales lines from database (in case of they exist) caused by parallel processing
 
    Args:
        Nothing
    Returns: 
        Nothing
    """
    pipeline = [
        {"$group": {"_id": {"file": "$file", "line": "$line"}, "unique_ids": {"$push": "$_id"}, "clean_ids": {"$push": "$clean"}, "count": {"$sum": 1}}},
        {"$match": {"count": { "$gte": 2 }}}
    ]
    cursor = collection.aggregate(pipeline, allowDiskUse = True)
    response = []
    for doc in cursor:
        response.append(doc["unique_ids"][1:])
    to_del = [item for sublist in response for item in sublist]
 
    flat_list = [DeleteOne({"_id":i}) for i in to_del]
    collection.bulk_write(flat_list, ordered=False)
    
 
def upload_main(event, s3_session):
    """ Orchestra every function and executes find_irsc function using multiprocessing 
    
    Args:
        event (dict): Dictionary result of ProcessFiles.py
        s3_session: (boto3.session object): S3 session to connect to AWS bucket
    Returns: Nothing
    """
    cat_match = event["cat_match"]
    tag = event["tag"]
    all_files = []
    s3_filename = "s3://"+event["results_bucket"]+"/"+event["results_path"]+"/"
    for i in event["format"]:
        q = delete_files_mongoDB(event["format"][i]["files"])
        for j in q:
            j = j+".parquet"
            all_files.append(s3_filename+j)
    dfs = wr.s3.read_parquet(all_files, boto3_session=s3_session, chunked=500_000)
    batch = 1
    for df in dfs:
        # df["quantity"] = df["quantity"].fillna(0)
        # df["total_local"] = df["total_local"].fillna(0.0)
        # df["full_total_foreign"] = df["full_total_foreign"].fillna(0.0)
        # df["quantity"] = df["quantity"].astype("int")
        # df["total_local"] = df["total_local"].astype("float")
        # df["full_total_foreign"] = df["full_total_foreign"].astype("float")
        # df.fillna("undefined", inplace=True)
        # df.reset_index(drop=True, inplace=True)
        # df["catalogue_id"] = "no_catalogue_id"
        if str(cat_match) == "True" or str(cat_match) == "CatType.true": # Launches catalogue matching
            print(f"{'*'*15} Catalogue Matching Starting {'*'*15}")
            df["clean"] = 0
            mp.freeze_support()
            isrcs = list(df["isrc_id"].unique())
            pool = mp.Pool(processes=6)
            pool.starmap(find_irsc, zip(isrcs, repeat(df)))
            pool.close()
            pool.join()
        else: # Just uploads sales lines
            print(f"{'*'*15} Uploading Data to Mongo {'*'*15}")
            upload_data_mongo(df)
        print(f"Batch: {batch} finished. Lines uploaded to {collection_name}")
        batch+=1
    # try:
    #     delete_duplicates()
    # except Exception:
    #     print("No files duplicated")
    #     pass