<commit_before># Copyright 2015 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you
# may not use this file except in compliance with the License.  You
# may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.  See the License for the specific language governing
# permissions and limitations under the License.

"""
MongoDB store for Coal Mine
"""

from coal_mine.abstract_store import AbstractStore
import bson
from copy import copy
from logbook import Logger
from pymongo import MongoClient, IndexModel, ASCENDING
from pymongo.errors import AutoReconnect
import re
import ssl
import time

log = Logger('MongoStore')


class MongoStore(AbstractStore):
    def __init__(self, hosts, database, username, password, **kwargs):
        """Keyword arguments are the same as what would be passed to
        MongoClient."""

        if 'ssl_cert_reqs' in kwargs:
            if kwargs['ssl_cert_reqs'] == 'NONE':
                kwargs['ssl_cert_reqs'] = ssl.CERT_NONE
            elif kwargs['ssl_cert_reqs'] == 'OPTIONAL':
                kwargs['ssl_cert_reqs'] = ssl.CERT_OPTIONAL
            elif kwargs['ssl_cert_reqs'] == 'REQUIRED':
                kwargs['ssl_cert_reqs'] = ssl.CERT_REQUIRED
            else:
                raise TypeError('ssl_cert_reqs must be NONE, OPTIONAL, or '
                                'REQUIRED')
        connection = MongoClient(hosts, **kwargs)
        db = connection[database]
        if username or password:
            db.authenticate(username, password)
        self.db = db
        self.collection = self.db['canaries']

        self.collection.create_indexes([
            IndexModel([('id', ASCENDING)]),
            # for list()
            IndexModel([('paused', ASCENDING),
                        ('late', ASCENDING),
                        ('deadline', ASCENDING)]),
            IndexModel([('paused', ASCENDING), ('deadline', ASCENDING)]),
            IndexModel([('late', ASCENDING), ('deadline', ASCENDING)]),
            # for find_identifier(), as well as to ensure unique slugs
            IndexModel([('slug', ASCENDING)], unique=True)])

    def create(self, canary):
        canary['_id'] = bson.ObjectId()
        while True:
            try:
                self.collection.insert_one(canary)
                del canary['_id']
                break
            except AutoReconnect:
                log.exception('save failure, retrying')
                time.sleep(1)

    def update(self, identifier, updates):<commit_msg>BUG: 'id' index is supposed to be unique<commit_after># Copyright 2015 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you
# may not use this file except in compliance with the License.  You
# may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.  See the License for the specific language governing
# permissions and limitations under the License.

"""
MongoDB store for Coal Mine
"""

from coal_mine.abstract_store import AbstractStore
import bson
from copy import copy
from logbook import Logger
from pymongo import MongoClient, IndexModel, ASCENDING
from pymongo.errors import AutoReconnect
import re
import ssl
import time

log = Logger('MongoStore')


class MongoStore(AbstractStore):
    def __init__(self, hosts, database, username, password, **kwargs):
        """Keyword arguments are the same as what would be passed to
        MongoClient."""

        if 'ssl_cert_reqs' in kwargs:
            if kwargs['ssl_cert_reqs'] == 'NONE':
                kwargs['ssl_cert_reqs'] = ssl.CERT_NONE
            elif kwargs['ssl_cert_reqs'] == 'OPTIONAL':
                kwargs['ssl_cert_reqs'] = ssl.CERT_OPTIONAL
            elif kwargs['ssl_cert_reqs'] == 'REQUIRED':
                kwargs['ssl_cert_reqs'] = ssl.CERT_REQUIRED
            else:
                raise TypeError('ssl_cert_reqs must be NONE, OPTIONAL, or '
                                'REQUIRED')
        connection = MongoClient(hosts, **kwargs)
        db = connection[database]
        if username or password:
            db.authenticate(username, password)
        self.db = db
        self.collection = self.db['canaries']

        # Compatibility for the fact that there were a couple releases (0.4.1
        # through 0.4.3) when the the code was creating the index on id
        # without unique=True.
        existing_indexes = self.collection.index_information()
        try:
            id_index = existing_indexes['id_1']
            try:
                is_unique = id_index['unique']
            except:
                is_unique = False
            if not is_unique:
                self.collection.drop_index('id_1')
        except:
            pass

        self.collection.create_indexes([
            IndexModel([('id', ASCENDING)], unique=True),
            # for list()
            IndexModel([('paused', ASCENDING),
                        ('late', ASCENDING),
                        ('deadline', ASCENDING)]),
            IndexModel([('paused', ASCENDING), ('deadline', ASCENDING)]),
            IndexModel([('late', ASCENDING), ('deadline', ASCENDING)]),
            # for find_identifier(), as well as to ensure unique slugs
            IndexModel([('slug', ASCENDING)], unique=True)])

    def create(self, canary):
        canary['_id'] = bson.ObjectId()
        while True:
            try:
                self.collection.insert_one(canary)
                del canary['_id']
                break
            except AutoReconnect:
                log.exception('save failure, retrying')
                time.sleep(1)

    def update(self, identifier, updates):
{'fraction_non_alphanumeric': 0.09911822338340953, 'fraction_numerical': 0.004245591116917048, 'mean_word_length': 2.3506564551422318, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2065256', 'n_tokens_mistral': 1711, 'n_tokens_neox': 1644, 'n_words': 523}