mirror of
https://github.com/samuelclay/NewsBlur.git
synced 2025-08-05 16:58:59 +00:00
555 lines
20 KiB
Python
555 lines
20 KiB
Python
import re
|
|
import time
|
|
import datetime
|
|
import pymongo
|
|
import pyelasticsearch
|
|
import redis
|
|
import celery
|
|
import html
|
|
import mongoengine as mongo
|
|
from django.conf import settings
|
|
from django.contrib.auth.models import User
|
|
from apps.search.tasks import IndexSubscriptionsForSearch
|
|
from apps.search.tasks import FinishIndexSubscriptionsForSearch
|
|
from apps.search.tasks import IndexSubscriptionsChunkForSearch
|
|
from apps.search.tasks import IndexFeedsForSearch
|
|
from utils import log as logging
|
|
from utils.feed_functions import chunks
|
|
|
|
class MUserSearch(mongo.Document):
|
|
'''Search index state of a user's subscriptions.'''
|
|
user_id = mongo.IntField(unique=True)
|
|
last_search_date = mongo.DateTimeField()
|
|
subscriptions_indexed = mongo.BooleanField()
|
|
subscriptions_indexing = mongo.BooleanField()
|
|
|
|
meta = {
|
|
'collection': 'user_search',
|
|
'indexes': ['user_id'],
|
|
'allow_inheritance': False,
|
|
}
|
|
|
|
@classmethod
|
|
def get_user(cls, user_id, create=True):
|
|
try:
|
|
user_search = cls.objects.read_preference(pymongo.ReadPreference.PRIMARY)\
|
|
.get(user_id=user_id)
|
|
except cls.DoesNotExist:
|
|
if create:
|
|
user_search = cls.objects.create(user_id=user_id)
|
|
else:
|
|
user_search = None
|
|
|
|
return user_search
|
|
|
|
def touch_search_date(self):
|
|
if not self.subscriptions_indexed and not self.subscriptions_indexing:
|
|
self.schedule_index_subscriptions_for_search()
|
|
self.subscriptions_indexing = True
|
|
|
|
self.last_search_date = datetime.datetime.now()
|
|
self.save()
|
|
|
|
def schedule_index_subscriptions_for_search(self):
|
|
IndexSubscriptionsForSearch.apply_async(kwargs=dict(user_id=self.user_id),
|
|
queue='search_indexer')
|
|
|
|
# Should be run as a background task
|
|
def index_subscriptions_for_search(self):
|
|
from apps.rss_feeds.models import Feed
|
|
from apps.reader.models import UserSubscription
|
|
|
|
SearchStory.create_elasticsearch_mapping()
|
|
|
|
start = time.time()
|
|
user = User.objects.get(pk=self.user_id)
|
|
r = redis.Redis(connection_pool=settings.REDIS_PUBSUB_POOL)
|
|
r.publish(user.username, 'search_index_complete:start')
|
|
|
|
subscriptions = UserSubscription.objects.filter(user=user).only('feed')
|
|
total = subscriptions.count()
|
|
|
|
feed_ids = []
|
|
for sub in subscriptions:
|
|
try:
|
|
feed_ids.append(sub.feed.pk)
|
|
except Feed.DoesNotExist:
|
|
continue
|
|
|
|
feed_id_chunks = [c for c in chunks(feed_ids, 6)]
|
|
logging.user(user, "~FCIndexing ~SB%s feeds~SN in %s chunks..." %
|
|
(total, len(feed_id_chunks)))
|
|
|
|
search_chunks = [IndexSubscriptionsChunkForSearch.s(feed_ids=feed_id_chunk,
|
|
user_id=self.user_id
|
|
).set(queue='search_indexer')
|
|
for feed_id_chunk in feed_id_chunks]
|
|
callback = FinishIndexSubscriptionsForSearch.s(user_id=self.user_id,
|
|
start=start).set(queue='search_indexer')
|
|
celery.chord(search_chunks)(callback)
|
|
|
|
def finish_index_subscriptions_for_search(self, start):
|
|
from apps.reader.models import UserSubscription
|
|
|
|
r = redis.Redis(connection_pool=settings.REDIS_PUBSUB_POOL)
|
|
user = User.objects.get(pk=self.user_id)
|
|
subscriptions = UserSubscription.objects.filter(user=user).only('feed')
|
|
total = subscriptions.count()
|
|
duration = time.time() - start
|
|
|
|
logging.user(user, "~FCIndexed ~SB%s feeds~SN in ~FM~SB%s~FC~SN sec." %
|
|
(total, round(duration, 2)))
|
|
r.publish(user.username, 'search_index_complete:done')
|
|
|
|
self.subscriptions_indexed = True
|
|
self.subscriptions_indexing = False
|
|
self.save()
|
|
|
|
def index_subscriptions_chunk_for_search(self, feed_ids):
|
|
from apps.rss_feeds.models import Feed
|
|
r = redis.Redis(connection_pool=settings.REDIS_PUBSUB_POOL)
|
|
user = User.objects.get(pk=self.user_id)
|
|
|
|
logging.user(user, "~FCIndexing %s feeds..." % len(feed_ids))
|
|
|
|
for feed_id in feed_ids:
|
|
feed = Feed.get_by_id(feed_id)
|
|
if not feed: continue
|
|
|
|
feed.index_stories_for_search()
|
|
|
|
r.publish(user.username, 'search_index_complete:feeds:%s' %
|
|
','.join([str(f) for f in feed_ids]))
|
|
|
|
@classmethod
|
|
def schedule_index_feeds_for_search(cls, feed_ids, user_id):
|
|
user_search = cls.get_user(user_id, create=False)
|
|
if (not user_search or
|
|
not user_search.subscriptions_indexed or
|
|
user_search.subscriptions_indexing):
|
|
# User hasn't searched before.
|
|
return
|
|
|
|
if not isinstance(feed_ids, list):
|
|
feed_ids = [feed_ids]
|
|
IndexFeedsForSearch.apply_async(kwargs=dict(feed_ids=feed_ids, user_id=user_id),
|
|
queue='search_indexer')
|
|
|
|
@classmethod
|
|
def index_feeds_for_search(cls, feed_ids, user_id):
|
|
from apps.rss_feeds.models import Feed
|
|
user = User.objects.get(pk=user_id)
|
|
|
|
logging.user(user, "~SB~FCIndexing %s~FC by request..." % feed_ids)
|
|
|
|
for feed_id in feed_ids:
|
|
feed = Feed.get_by_id(feed_id)
|
|
if not feed: continue
|
|
|
|
feed.index_stories_for_search()
|
|
|
|
@classmethod
|
|
def remove_all(cls, drop_index=False):
|
|
# You only need to drop the index if there is data you want to clear.
|
|
# A new search server won't need this, as there isn't anything to drop.
|
|
if drop_index:
|
|
logging.info(" ---> ~FRRemoving stories search index...")
|
|
SearchStory.drop()
|
|
|
|
user_searches = cls.objects.all()
|
|
logging.info(" ---> ~SN~FRRemoving ~SB%s~SN user searches..." % user_searches.count())
|
|
for user_search in user_searches:
|
|
try:
|
|
user_search.remove()
|
|
except Exception as e:
|
|
print(" ****> Error on search removal: %s" % e)
|
|
|
|
def remove(self):
|
|
from apps.rss_feeds.models import Feed
|
|
from apps.reader.models import UserSubscription
|
|
|
|
user = User.objects.get(pk=self.user_id)
|
|
subscriptions = UserSubscription.objects.filter(user=self.user_id)
|
|
total = subscriptions.count()
|
|
removed = 0
|
|
|
|
for sub in subscriptions:
|
|
try:
|
|
feed = sub.feed
|
|
except Feed.DoesNotExist:
|
|
continue
|
|
if not feed.search_indexed:
|
|
continue
|
|
feed.search_indexed = False
|
|
feed.save()
|
|
removed += 1
|
|
|
|
logging.user(user, "~FCRemoved ~SB%s/%s feed's search indexes~SN for ~SB~FB%s~FC~SN." %
|
|
(removed, total, user.username))
|
|
self.delete()
|
|
|
|
class SearchStory:
|
|
|
|
_es_client = None
|
|
name = "stories"
|
|
|
|
@classmethod
|
|
def ES(cls):
|
|
if cls._es_client is None:
|
|
cls._es_client = pyelasticsearch.ElasticSearch(settings.ELASTICSEARCH_STORY_HOST)
|
|
cls.create_elasticsearch_mapping()
|
|
return cls._es_client
|
|
|
|
@classmethod
|
|
def index_name(cls):
|
|
return "%s-index" % cls.name
|
|
|
|
@classmethod
|
|
def create_elasticsearch_mapping(cls, delete=False):
|
|
if delete:
|
|
cls.ES().delete_index(cls.index_name())
|
|
|
|
try:
|
|
cls.ES().create_index(cls.index_name())
|
|
logging.debug(" ---> ~FCCreating search index for ~FM%s" % cls.index_name())
|
|
except pyelasticsearch.IndexAlreadyExistsError:
|
|
return
|
|
|
|
mapping = {
|
|
'title': {
|
|
'boost': 3.0,
|
|
'store': False,
|
|
'type': 'text',
|
|
'analyzer': 'snowball',
|
|
},
|
|
'content': {
|
|
'boost': 1.0,
|
|
'store': False,
|
|
'type': 'text',
|
|
'analyzer': 'snowball',
|
|
},
|
|
'tags': {
|
|
'boost': 2.0,
|
|
'store': False,
|
|
'type': 'keyword',
|
|
},
|
|
'author': {
|
|
'boost': 1.0,
|
|
'store': False,
|
|
'type': 'text',
|
|
'analyzer': 'simple',
|
|
},
|
|
'feed_id': {
|
|
'store': False,
|
|
'type': 'integer'
|
|
},
|
|
'date': {
|
|
'store': False,
|
|
'type': 'date',
|
|
}
|
|
}
|
|
cls.ES().put_mapping(index=cls.index_name(), doc_type='story-type', mapping={
|
|
'properties': mapping,
|
|
})
|
|
cls.ES().flush(cls.index_name())
|
|
|
|
@classmethod
|
|
def index(cls, story_hash, story_title, story_content, story_tags, story_author, story_feed_id,
|
|
story_date):
|
|
cls.create_elasticsearch_mapping()
|
|
|
|
doc = {
|
|
"content" : story_content,
|
|
"title" : story_title,
|
|
"tags" : ', '.join(story_tags),
|
|
"author" : story_author,
|
|
"feed_id" : story_feed_id,
|
|
"date" : story_date,
|
|
}
|
|
try:
|
|
cls.ES().index(index=cls.index_name(), doc_type='story-type', doc=doc, id=story_hash)
|
|
except pyelasticsearch.ElasticHttpError as e:
|
|
logging.debug(" ***> ~FRNo search server available for story indexing.")
|
|
if settings.DEBUG:
|
|
raise e
|
|
|
|
@classmethod
|
|
def remove(cls, story_hash):
|
|
try:
|
|
cls.ES().delete(index=cls.index_name(), id=story_hash)
|
|
except pyelasticsearch.ElasticHttpError:
|
|
logging.debug(" ***> ~FRNo search server available for story deletion.")
|
|
|
|
@classmethod
|
|
def drop(cls):
|
|
try:
|
|
cls.ES().delete_index(cls.index_name())
|
|
except pyelasticsearch.ElasticHttpNotFoundError:
|
|
logging.debug(" ***> ~FBNo index found, nothing to drop.")
|
|
|
|
|
|
@classmethod
|
|
def query(cls, feed_ids, query, order, offset, limit, strip=False):
|
|
try:
|
|
cls.ES().flush(index=cls.index_name())
|
|
except pyelasticsearch.ElasticHttpError:
|
|
logging.debug(" ***> ~FRNo search server available.")
|
|
return []
|
|
|
|
if strip:
|
|
query = re.sub(r'([^\s\w_\-])+', ' ', query) # Strip non-alphanumeric
|
|
query = html.unescape(query)
|
|
|
|
body = {
|
|
"query": {
|
|
"bool": {
|
|
"must": [
|
|
{"query_string": { "query": query, "default_operator": "AND" }},
|
|
{"terms": { "feed_id": feed_ids[:2000] }},
|
|
]
|
|
}
|
|
},
|
|
'sort': [{'date': {'order': 'desc' if order == "newest" else "asc"}}],
|
|
'from': offset,
|
|
'size': limit
|
|
}
|
|
try:
|
|
results = cls.ES().search(body, index=cls.index_name())
|
|
except pyelasticsearch.ElasticHttpError as e:
|
|
logging.debug(" ***> ~FRNo search server available for querying: %s" % e)
|
|
return []
|
|
|
|
# s = elasticsearch_dsl.Search(using=cls.ES(), index=cls.index_name())
|
|
# string_q = elasticsearch_dsl.Q('query_string', query=query, default_operator="AND")
|
|
# feed_q = elasticsearch_dsl.Q('terms', feed_id=feed_ids[:2000])
|
|
# search_q = string_q & feed_q
|
|
# s = s.query(search_q)
|
|
# s = s.sort(sort)[offset:offset+limit]
|
|
# results = s.execute()
|
|
|
|
# string_q = pyes.query.QueryStringQuery(query, default_operator="AND")
|
|
# feed_q = pyes.query.TermsQuery('feed_id', feed_ids[:2000])
|
|
# q = pyes.query.BoolQuery(must=[string_q, feed_q])
|
|
# try:
|
|
# results = cls.ES().search(q, indices=cls.index_name(),
|
|
# partial_fields={}, sort=sort, start=offset, size=limit)
|
|
# except elasticsearch.exceptions.ConnectionError:
|
|
# logging.debug(" ***> ~FRNo search server available.")
|
|
# return []
|
|
|
|
logging.info(" ---> ~FG~SNSearch ~FCstories~FG for: ~SB%s~SN, ~SB%s~SN results (across %s feed%s)" %
|
|
(query, len(results['hits']['hits']), len(feed_ids), 's' if len(feed_ids) != 1 else ''))
|
|
|
|
try:
|
|
result_ids = [r['_id'] for r in results['hits']['hits']]
|
|
except Exception as e:
|
|
logging.info(" ---> ~FRInvalid search query \"%s\": %s" % (query, e))
|
|
return []
|
|
|
|
return result_ids
|
|
|
|
@classmethod
|
|
def global_query(cls, query, order, offset, limit, strip=False):
|
|
cls.create_elasticsearch_mapping()
|
|
cls.ES().refresh()
|
|
|
|
if strip:
|
|
query = re.sub(r'([^\s\w_\-])+', ' ', query) # Strip non-alphanumeric
|
|
query = html.unescape(query)
|
|
|
|
body = {
|
|
"query": {
|
|
"bool": {
|
|
"must": [
|
|
{"query_string": { "query": query, "default_operator": "AND" }},
|
|
]
|
|
}
|
|
},
|
|
'sort': [{'date': {'order': 'desc' if order == "newest" else "asc"}}],
|
|
'from': offset,
|
|
'size': limit
|
|
}
|
|
try:
|
|
results = cls.ES().search(body, index=cls.index_name())
|
|
except pyelasticsearch.ElasticHttpError as e:
|
|
logging.debug(" ***> ~FRNo search server available for querying: %s" % e)
|
|
return []
|
|
|
|
# sort = "date:desc" if order == "newest" else "date:asc"
|
|
# string_q = pyes.query.QueryStringQuery(query, default_operator="AND")
|
|
# try:
|
|
# results = cls.ES().search(string_q, indices=cls.index_name(),
|
|
# partial_fields={}, sort=sort, start=offset, size=limit)
|
|
# except elasticsearch.exceptions.ConnectionError:
|
|
# logging.debug(" ***> ~FRNo search server available.")
|
|
# return []
|
|
|
|
logging.info(" ---> ~FG~SNSearch ~FCstories~FG for: ~SB%s~SN (across all feeds)" %
|
|
(query))
|
|
|
|
try:
|
|
result_ids = [r['_id'] for r in results['hits']['hits']]
|
|
except Exception as e:
|
|
logging.info(" ---> ~FRInvalid search query \"%s\": %s" % (query, e))
|
|
return []
|
|
|
|
return result_ids
|
|
|
|
|
|
class SearchFeed:
|
|
|
|
_es_client = None
|
|
name = "feeds"
|
|
|
|
@classmethod
|
|
def ES(cls):
|
|
if cls._es_client is None:
|
|
cls._es_client = pyelasticsearch.ElasticSearch(settings.ELASTICSEARCH_FEED_HOST)
|
|
cls.create_elasticsearch_mapping()
|
|
return cls._es_client
|
|
|
|
@classmethod
|
|
def index_name(cls):
|
|
return "%s-index" % cls.name
|
|
|
|
@classmethod
|
|
def create_elasticsearch_mapping(cls, delete=False):
|
|
if delete:
|
|
logging.debug(" ---> ~FRDeleting search index for ~FM%s" % cls.index_name())
|
|
cls.ES().delete_index(cls.index_name())
|
|
|
|
try:
|
|
index_settings = {
|
|
"index" : {
|
|
"analysis": {
|
|
"analyzer": {
|
|
"edgengram_analyzer": {
|
|
"filter": ["edgengram_analyzer"],
|
|
"tokenizer": "lowercase",
|
|
"type": "custom"
|
|
},
|
|
},
|
|
"filter": {
|
|
"edgengram_analyzer": {
|
|
"max_gram": "15",
|
|
"min_gram": "1",
|
|
"type": "edge_ngram"
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
cls.ES().create_index(cls.index_name(), settings=index_settings)
|
|
logging.debug(" ---> ~FCCreating search index for ~FM%s" % cls.index_name())
|
|
except pyelasticsearch.IndexAlreadyExistsError:
|
|
return
|
|
|
|
mapping = {
|
|
"feed_address": {
|
|
'analyzer': 'snowball',
|
|
"store": False,
|
|
"term_vector": "with_positions_offsets",
|
|
"type": "text"
|
|
},
|
|
"feed_id": {
|
|
"store": True,
|
|
"type": "text"
|
|
},
|
|
"num_subscribers": {
|
|
"store": True,
|
|
"type": "long"
|
|
},
|
|
"title": {
|
|
"analyzer": "snowball",
|
|
"store": False,
|
|
"term_vector": "with_positions_offsets",
|
|
"type": "text"
|
|
},
|
|
"link": {
|
|
"analyzer": "snowball",
|
|
"store": False,
|
|
"term_vector": "with_positions_offsets",
|
|
"type": "text"
|
|
}
|
|
}
|
|
cls.ES().put_mapping(index=cls.index_name(), doc_type='feeds-type', mapping={
|
|
'properties': mapping,
|
|
})
|
|
cls.ES().flush(cls.index_name())
|
|
|
|
@classmethod
|
|
def index(cls, feed_id, title, address, link, num_subscribers):
|
|
doc = {
|
|
"feed_id": feed_id,
|
|
"title": title,
|
|
"feed_address": address,
|
|
"link": link,
|
|
"num_subscribers": num_subscribers,
|
|
}
|
|
try:
|
|
cls.ES().index(index=cls.index_name(), doc_type='feeds-type', doc=doc, id=feed_id)
|
|
except pyelasticsearch.ElasticHttpError:
|
|
logging.debug(" ***> ~FRNo search server available for feed indexing.")
|
|
|
|
@classmethod
|
|
def query(cls, text, max_subscribers=5):
|
|
try:
|
|
cls.ES().flush(index=cls.index_name())
|
|
except pyelasticsearch.ElasticHttpError:
|
|
logging.debug(" ***> ~FRNo search server available for feed querying.")
|
|
return []
|
|
|
|
if settings.DEBUG:
|
|
max_subscribers = 1
|
|
|
|
body = {
|
|
"query": {
|
|
"bool": {
|
|
"should": [
|
|
{"match": { "address": { "query": text, 'cutoff_frequency': "0.0005", 'minimum_should_match': "75%" } }},
|
|
{"match": { "title": { "query": text, 'cutoff_frequency': "0.0005", 'minimum_should_match': "75%" } }},
|
|
{"match": { "link": { "query": text, 'cutoff_frequency': "0.0005", 'minimum_should_match': "75%" } }},
|
|
]
|
|
}
|
|
},
|
|
'sort': [{'num_subscribers': {'order': 'desc'}}],
|
|
}
|
|
try:
|
|
results = cls.ES().search(body, doc_type='feeds-type', index=cls.index_name())
|
|
except pyelasticsearch.ElasticHttpError as e:
|
|
logging.debug(" ***> ~FRNo search server available for feed querying: %s" % e)
|
|
return []
|
|
|
|
# s = elasticsearch_dsl.Search(using=cls.ES(), index=cls.index_name())
|
|
# address = elasticsearch_dsl.Q('match', address=text)
|
|
# link = elasticsearch_dsl.Q('match', link=text)
|
|
# title = elasticsearch_dsl.Q('match', title=text)
|
|
# search_q = address | link | title
|
|
# s = s.query(search_q).extra(cutoff_frequency="0.0005", minimum_should_match="75%")
|
|
# s = s.sort("-num_subscribers")
|
|
# body = s.to_dict()
|
|
# print(f"Before: {body}")
|
|
# results = s.execute()
|
|
|
|
# q = pyes.query.BoolQuery()
|
|
# q.add_should(pyes.query.MatchQuery('address', text, analyzer="simple", cutoff_frequency=0.0005, minimum_should_match="75%"))
|
|
# q.add_should(pyes.query.MatchQuery('link', text, analyzer="simple", cutoff_frequency=0.0005, minimum_should_match="75%"))
|
|
# q.add_should(pyes.query.MatchQuery('title', text, analyzer="simple", cutoff_frequency=0.0005, minimum_should_match="75%"))
|
|
# q = pyes.Search(q, min_score=1)
|
|
# results = cls.ES().search(query=q, size=max_subscribers, sort="num_subscribers:desc")
|
|
|
|
logging.info("~FGSearch ~FCfeeds~FG: ~SB%s~SN, ~SB%s~SN results" % (text, len(results['hits']['hits'])))
|
|
|
|
return results['hits']['hits']
|
|
|
|
@classmethod
|
|
def export_csv(cls):
|
|
import djqscsv
|
|
from apps.rss_feeds.models import Feed
|
|
|
|
qs = Feed.objects.filter(num_subscribers__gte=20).values('id', 'feed_title', 'feed_address', 'feed_link', 'num_subscribers')
|
|
csv = djqscsv.render_to_csv_response(qs).content
|
|
f = open('feeds.csv', 'w+')
|
|
f.write(csv)
|
|
f.close()
|
|
|