<commit_before>"""
Universal database functions.
specific functions related to some components are in component file
"""

import random
import urllib.parse
from datetime import datetime
from random import shuffle

import pymongo
from upol_crawler.settings import *
from upol_crawler.utils import urls


def init(db):
    """Database init, create indexes"""
    db['Urls'].create_index('visited')
    db['Urls'].create_index('queued')
    db['Urls'].create_index('timeout')
    db['Urls'].create_index('content.hashes.document')
    db['Limiter'].create_index('ip', unique=True)


def _prepare_url_object(url, visited, queued, depth):
    """Prepare url object before inserting into database"""
    url_object = {'_id': urls.hash(url),
                  'url': url,
                  'domain': urls.domain(url),
                  'depth': depth,
                  'visited': visited,
                  'queued': queued,
                  'progress': {'discovered': str(datetime.now())}}

    return url_object


def _prepare_file_url_object(url, response, depth):
    """Prepare file url object before inserting into database"""
    # Prepare content_type, remove charset etc
    # for example 'text/html; charset=utf-8'
    content_type = response.headers.get('Content-Type')

    if content_type is not None:
        content_type = content_type.split(';')[0]

    file_url_object = {'_id': urls.hash(url),
                       'url': url,
                       'domain': urls.domain(url),
                       'depth': depth,
                       'content_type': content_type,
                       'content_length': response.headers.get('Content-Length'),
                       'progress': {'discovered': str(datetime.now())}}

    return file_url_object


def insert_url(db, url, visited, queued, depth):
    """Insert url into db"""
    url_object = _prepare_url_object(url, visited, queued, depth)

    try:
        result = db['Urls'].insert_one(url_object).inserted_id
    except pymongo.errors.DuplicateKeyError as e:
        return False

    return result


def inser_file_url(db, url, response, depth):
    """Insert file url into db"""
    file_url_object = _prepare_file_url_object(url, response, depth)

    try:
        result = db['Files'].insert_one(file_url_object).inserted_id
    except pymongo.errors.DuplicateKeyError as e:
        return False

    return result


def batch_insert_url(db, urls_with_depths, visited, queued):
    """Inser batch of urls into db"""

    url_documents = []

    for url in urls_with_depths:
        url_object = _prepare_url_object(url.get('url'),
                                         visited,
                                         queued,
                                         url.get('depth'))
        url_documents.append(url_object)

    try:
        result = db['Urls'].insert_many(url_documents, ordered=False)
    except pymongo.errors.BulkWriteError:
        # TODO - There is no point of returning result variable from this function. insert_many can fail on one url because of duplicity and thats totally fine. So probably better to ignore return statement
        result = None

    return result


def insert_url_info(db, url, info_type, arg={}):
    """Insert aditional info about url into database"""
    collection = db[info_type]

    log_object = {'_id': urls.hash(url),
                  'url': url}

    for key, depth in arg.items():
        log_object[key] = str(depth)

    try:
        collection.insert_one(log_object).inserted_id
    except pymongo.errors.DuplicateKeyError as e:
        return False


def delete_url(db, url):
    """Try to delete url from db, returns True if case of success"""
    result = db['Urls'].delete_one({'_id': urls.hash(url)})

    return result.deleted_count > 0


def set_visited_url(db, url, response, content):
    """Try to set url to visited and update other important informations"""
    url_hash = urls.hash(url)

    is_permanent_redirect = False

    for history in response.history:
        if history.is_permanent_redirect:
            is_permanent_redirect = True
            break<commit_msg>get_or_create_canonical_group function added<commit_after>"""
Universal database functions.
specific functions related to some components are in component file
"""

import random
import urllib.parse
from datetime import datetime
from random import shuffle

import pymongo
from upol_crawler.settings import *
from upol_crawler.utils import urls


def init(db):
    """Database init, create indexes"""
    db['Urls'].create_index('visited')
    db['Urls'].create_index('queued')
    db['Urls'].create_index('timeout')
    db['Urls'].create_index('content.hashes.document')
    db['Limiter'].create_index('ip', unique=True)


def _prepare_url_object(url, visited, queued, depth):
    """Prepare url object before inserting into database"""
    url_object = {'_id': urls.hash(url),
                  'url': url,
                  'domain': urls.domain(url),
                  'depth': depth,
                  'visited': visited,
                  'queued': queued,
                  'progress': {'discovered': str(datetime.now())}}

    return url_object


def _prepare_file_url_object(url, response, depth):
    """Prepare file url object before inserting into database"""
    # Prepare content_type, remove charset etc
    # for example 'text/html; charset=utf-8'
    content_type = response.headers.get('Content-Type')

    if content_type is not None:
        content_type = content_type.split(';')[0]

    file_url_object = {'_id': urls.hash(url),
                       'url': url,
                       'domain': urls.domain(url),
                       'depth': depth,
                       'content_type': content_type,
                       'content_length': response.headers.get('Content-Length'),
                       'progress': {'discovered': str(datetime.now())}}

    return file_url_object


def insert_url(db, url, visited, queued, depth):
    """Insert url into db"""
    url_object = _prepare_url_object(url, visited, queued, depth)

    try:
        result = db['Urls'].insert_one(url_object).inserted_id
    except pymongo.errors.DuplicateKeyError as e:
        return False

    return result


def inser_file_url(db, url, response, depth):
    """Insert file url into db"""
    file_url_object = _prepare_file_url_object(url, response, depth)

    try:
        result = db['Files'].insert_one(file_url_object).inserted_id
    except pymongo.errors.DuplicateKeyError as e:
        return False

    return result


def batch_insert_url(db, urls_with_depths, visited, queued):
    """Inser batch of urls into db"""

    url_documents = []

    for url in urls_with_depths:
        url_object = _prepare_url_object(url.get('url'),
                                         visited,
                                         queued,
                                         url.get('depth'))
        url_documents.append(url_object)

    try:
        result = db['Urls'].insert_many(url_documents, ordered=False)
    except pymongo.errors.BulkWriteError:
        # TODO - There is no point of returning result variable from this function. insert_many can fail on one url because of duplicity and thats totally fine. So probably better to ignore return statement
        result = None

    return result


def insert_url_info(db, url, info_type, arg={}):
    """Insert aditional info about url into database"""
    collection = db[info_type]

    log_object = {'_id': urls.hash(url),
                  'url': url}

    for key, depth in arg.items():
        log_object[key] = str(depth)

    try:
        collection.insert_one(log_object).inserted_id
    except pymongo.errors.DuplicateKeyError as e:
        return False


def delete_url(db, url):
    """Try to delete url from db, returns True if case of success"""
    result = db['Urls'].delete_one({'_id': urls.hash(url)})

    return result.deleted_count > 0


def get_or_create_canonical_group(db, content_hash):
    """Try to get canonical group with given hash.
       Create new canonical group in case of fail."""

    # Possible chance of optimalization here
    canonical_group = list(db['CanonicalGroups'].find({'content_hash': content_hash}).limit(1))

    # Create new one
    if len(canonical_group) == 0:
        return db['CanonicalGroups'].insert({'content_hash': content_hash})
    else:
        return canonical_group[0].get('_id')


def set_visited_url(db, url, response, content):
    """Try to set url to visited and update other important informations"""
    url_hash = urls.hash(url)

    is_permanent_redirect = False

    for history in response.history:
        if history.is_permanent_redirect:
            is_permanent_redirect = True
            break
{'fraction_non_alphanumeric': 0.10619166856362394, 'fraction_numerical': 0.001024356931481903, 'mean_word_length': 2.668893528183716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '874051', 'n_tokens_mistral': 2571, 'n_tokens_neox': 2468, 'n_words': 688}