<commit_before>from __future__ import absolute_import

import json
import logging
import functools
from datetime import datetime

import requests
import cqlengine
from cqlengine import columns
from cqlengine import management
from cassandra.cluster import NoHostAvailable

from scrapi import settings


logger = logging.getLogger(__name__)


try:
    cqlengine.connection.setup(settings.CASSANDRA_URI, settings.CASSANDRA_KEYSPACE)
    management.create_keyspace(settings.CASSANDRA_KEYSPACE, replication_factor=1, strategy_class='SimpleStrategy')
except NoHostAvailable:
    logger.error('Could not connect to Cassandra, expect errors.')
    if settings.RECORD_HTTP_TRANSACTIONS:
        raise


class HarvesterResponse(cqlengine.Model):
    __table_name__ = 'responses'
    __keyspace__ = settings.CASSANDRA_KEYSPACE

    method = columns.Text(primary_key=True)
    url = columns.Text(primary_key=True, required=True, index=True)

    # Raw request data
    content = columns.Bytes()
    headers_str = columns.Text()
    time_made = columns.DateTime(default=datetime.now)

    @property
    def json(self):
        return json.loads(self.content)

    @property
    def headers(self):
        # TODO: make case insensitive multidict
        return json.loads(self.headers_str)


def record_or_load_response(method, url, **kwargs):
    try:
        return HarvesterResponse.get(url=url, method=method)
    except HarvesterResponse.DoesNotExist:
        response = requests.request(method, url, **kwargs)
        return HarvesterResponse(
            url=url,
            method=method,
            content=response.content
        ).save()


def request(method, url, **kwargs):
    if settings.RECORD_HTTP_TRANSACTIONS:
        return record_or_load_response(method, url)
    return requests.request(method, url, **kwargs)


# This has to be done after HarvesterResponse definition
management.sync_table(HarvesterResponse)

get = functools.partial(request, 'get')
put = functools.partial(request, 'put')
post = functools.partial(request, 'post')
delete = functools.partial(request, 'delete')<commit_msg>Clean up and store http response code as well<commit_after>from __future__ import absolute_import

import json
import logging
import functools
from datetime import datetime

import requests
import cqlengine
from cqlengine import columns

from scrapi import database  # noqa
from scrapi import settings


logger = logging.getLogger(__name__)


class HarvesterResponse(cqlengine.Model):
    __table_name__ = 'responses'
    __keyspace__ = settings.CASSANDRA_KEYSPACE

    method = columns.Text(primary_key=True)
    url = columns.Text(primary_key=True, required=True, index=True)

    # Raw request data
    content = columns.Bytes()
    headers_str = columns.Text()
    status_code = columns.Integer()
    time_made = columns.DateTime(default=datetime.now)

    @property
    def json(self):
        return json.loads(self.content)

    @property
    def headers(self):
        # TODO: make case insensitive multidict
        return json.loads(self.headers_str)


def record_or_load_response(method, url, **kwargs):
    try:
        return HarvesterResponse.get(url=url, method=method)
    except HarvesterResponse.DoesNotExist:
        response = requests.request(method, url, **kwargs)
        return HarvesterResponse(
            url=url,
            method=method,
            content=response.content,
            status_code=response.status_code,
            headers_str=json.dumps(response.headers)
        ).save()


def request(method, url, **kwargs):
    if settings.RECORD_HTTP_TRANSACTIONS:
        return record_or_load_response(method, url, **kwargs)
    return requests.request(method, url, **kwargs)


get = functools.partial(request, 'get')
put = functools.partial(request, 'put')
post = functools.partial(request, 'post')
delete = functools.partial(request, 'delete')
{'fraction_non_alphanumeric': 0.08747105737072293, 'fraction_numerical': 0.0002572678157962439, 'mean_word_length': 3.9591836734693877, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '2803625', 'n_tokens_mistral': 1147, 'n_tokens_neox': 1116, 'n_words': 270}