alexis2/db.py

from pymongo import MongoClient, errors, DESCENDING, ASCENDING
from bson.objectid import ObjectId
from datetime import timedelta

import datetime
import re
import news
import pytz
import time


class db():
    DB = "alexis"

    def __init__(self):
        self.client = MongoClient('localhost', 27017)
        self.c_sources = self.client[self.DB]["sources"]
        self.c_articles = self.client[self.DB]["articles"]
        self.c_kw = self.client[self.DB]["keywords"]

    def change_timezone(self, el):
        #gmt = pytz.timezone("America/Buenos_Aires")  # FIXME
        local = pytz.timezone("US/Central")
        est = pytz.timezone('US/Eastern')
        fmt = '%Y-%m-%d %H:%M:%S %Z%z'

        el["scrap_date"] = local.localize(el["scrap_date"])
        el["scrap_date"] = el["scrap_date"].astimezone(est)
        el["scrap_date"] = el["scrap_date"].strftime(fmt)
        return el

    def jsonable(self, el):
        if el is None or "_id" not in el:
            return el

        el["_id"] = str(el["_id"])
        return el

    def filter_articles(self, filter, limit=80):
        start = time.time()
        db_filter = {}
        db_filter["matches"] = True
        db_filter["deleted"] = False
        db_filter["scrap_date"] = {}
        if "maxdate" in filter and len(filter["maxdate"]) > 7:
            db_filter["scrap_date"]["$lt"] = datetime.datetime.strptime(
                filter["maxdate"], "%Y-%m-%d")

        if "mindate" in filter and len(filter["mindate"]) > 7:
            db_filter["scrap_date"]["$gt"] = datetime.datetime.strptime(
                filter["mindate"], "%Y-%m-%d")

        if db_filter["scrap_date"] == {}:
            del db_filter["scrap_date"]

        if "title" in filter and filter["title"] != "":
            filter["title"] = filter["title"].replace(
                "(", "\(").replace(")", "\)")
            db_filter["title"] = {"$regex": ".*%s.*" %
                                  filter["title"], "$options": 'i'}

        if "site" in filter and filter["site"] != "":
            filter["site"] = filter["site"].replace(
                "(", "\(").replace(")", "\)")
            db_filter["source_name"] = {
                "$regex": ".*%s.*" % filter["site"], "$options": 'i'}
            # db_filter["$or"]=[
            #    {"source_name": {"$regex":".*%s.*"%filter["site"], "$options": 'i'}},
            #    {"url": {"$regex":".*%s.*"%filter["site"], "$options": 'i' }}
            #    ]

        if "minweight" not in filter or not filter["minweight"].isdigit():
            filter["minweight"] = 1

        if "category" not in filter or filter["category"] == "-1":
            filter["category"] = -1

        if int(filter["category"]) >= 0:
            db_filter["category"] = int(filter["category"])
            # pass

        db_filter["weight"] = {"$gt": int(filter["minweight"])}

        if "keyword" in filter and len(filter["keyword"]) > 0:
            db_filter["kw"] = {"$all": filter["keyword"]}

        # db_filter["kw.0"]={"$exists":False}
        pageNumber = 0
        if "page" in filter:
            pageNumber = filter["page"]
        # print(db_filter)
        print(db_filter)
        fields = {"html": 0, "text": 0, "authors":0,"lowerurl":0}
        fields = {"kw":1,"scrap_date":1, "source_name":1,"summary":1,"title":1,"url":1,"weight":1}
        res = self.c_articles.find(db_filter, fields)
        if res is not None:
            toSkip = (((pageNumber - 1) * limit) if pageNumber > 0 else 0)
            #q = res.skip(toSkip).sort("scrap_date", DESCENDING).limit(limit)
            q = res.skip(toSkip).sort("scrap_date", DESCENDING).limit(limit*3) #+2 pages
            start=time.time()
        else:
            q = []
            count = 0
        ret_arts = [self.jsonable(self.change_timezone(p)) for p in q]
        count = len(ret_arts)+toSkip
        ret_arts = ret_arts[:limit]
        #count = len(ret_arts)
        print("pre-ret", time.time()-start)
        return {'articles': ret_arts,
                'count': count}

    def article(self, id):
        ret = self.c_articles.find_one({"_id": ObjectId(id)})
        return self.jsonable(ret)

    def delete_article(self, artid):
        return self.c_articles.update({"_id": ObjectId(artid)}, {"$set": {"deleted": True}})

    def undelete_article(self, artid):
        return self.c_articles.update({"_id": ObjectId(artid)}, {"$set": {"deleted": False}})

    def sources(self, id=None):
        if id is None:
            return [self.jsonable(p) for p in self.c_sources.find()]

        ret = self.c_sources.find_one({"_id": ObjectId(id)})
        return self.jsonable(ret)

    def rss_sources(self):
        return [p for p in self.c_sources.find({"rss": 1})]

    def insert_article(self, o):
        try:
            self.c_articles.insert_one(o)
            return True
        except errors.DuplicateKeyError:
            print("Dup")
            return False

    def count_error(self):
        return self.c_articles.count({'error': True})

    def purge_error(self):
        self.c_articles.remove({'error': True})

    def article_exists(self, url, src):
	#slow as balls
        #ret = self.c_articles.find({'source': src, 'lowerurl': url.lower()},{"_id":1}).limit(1) #faster than findone
        #articles may have different src and same url.

        ret = self.c_articles.find({'lowerurl': url.lower()},{"_id":1}).limit(1) #faster than findone
        c = len([ 1 for r in ret ] )
	#findOne returns the elem. find returns the cursor (pointer)
        return c>0

    def lower_arr(self, arr):
        return [a.lower() for a in arr]

    def to_regex(self, arr):  # "ACH" |  "PAYMENT" | "EBAY"
        return [re.compile('\\b' + w.lower() + '\\b') for w in arr]

    def get_keywords_list(self):
        kw = self.c_kw
        return [k["keyword"] for k in kw.find({})]

    def get_keywords_weight_list(self):
        kw = self.c_kw
        return list(kw.find({}, {'_id': False}).sort({'weight': 1}))

    def get_keywords(self, category):
        kw = self.c_kw
        ret = [{"weight": k["weight"],
                "kw": k["keyword"],
                "pattern": self.to_regex(k["keyword"].split("|"))} for k in kw.find({})]
        # FIXME, categories
        return ret

    def add_keyword(self, kw, weight):
        kwdb = self.c_kw
        try:
            data = {'keyword': kw, 'weight': weight}
            kwdb.insert_one(data)
            return True
        except errors.DuplicateKeyError:
            print("Dup KW")
            return False

    def update_keyword(self, kw):
        if 'oldkw' not in kw or 'weight' not in kw or 'newkw' not in kw:
            return False

        if not isinstance(kw['weight'], int):
            return False

        kwdb = self.c_kw
        try:
            kwdb.update({'keyword': kw['oldkw']},
                        {'$set':
                            {'weight': kw['weight'], 'keyword': kw['newkw']}})
            return True
        except Exception:
            print("Dup KW")
            return False

    def delete_keyword(self, kw):
        kwdb = self.c_kw
        try:
            kwdb.remove({'keyword': kw})
            return True
        except Exception:
            return False

    def add_source(self, url, rss, selector, name, category):
        self.c_sources.insert({"link": url,
                               "rss": rss,
                               "selector": selector,
                               "name": name,
                               "category": category})

    def reparse_articles(self, articles, kw):
        for a in articles:
            # print(a)
            art = self.c_articles.find_one({"url": a["url"]})
            res = news.match_keywords(kw, art["text"])
            #(matches,w_sum,matching)
            # if(res["matches"]):
            print(res)
            self.c_articles.update({"_id": ObjectId(a["_id"])}, {"$set": res})
            #matches, kw, weight

    def get_sources_with_error(self):
        r = self.c_articles.aggregate([{'$match': {'error': True}},
                                       {'$group': {'_id': '$source_name', 'count': {'$sum': 1}}}])
        return r

    def get_sources_with_count(self):
        # TODO filter with scrap_date, this year mb?
        r = self.c_articles.aggregate([{'$match': {'matches': True}},
                                       {'$group':
                                           {'_id': '$source_name',
                                            'count': {'$sum': 1},
                                            'weight': {'$avg': '$weight'}
                                            }
                                        }])

        # t = self.c_articles.aggregate([{'$group': {'_id': '$source_name', 'count': {'$sum': 1}}}])
        return r

    def get_keywords_with_count(self):
        this_year = datetime.datetime.now() - timedelta(days=180)
        r = self.c_articles.aggregate([{'$match': {
                                            '$and': [{'matches': True}, {'scrap_date': {'$gt': this_year}}]
                                        }},
                                       {'$unwind': '$kw'},
                                       {'$group':
                                           {'_id': '$kw', 'sum': {'$sum': 1}}
                                        }
                                       ])
        d1 = {}
        for res in r:
            d1[res['_id']] = res['sum']
        return d1