|
1
|
|
|
#!/usr/bin/env python |
|
2
|
|
|
# -*- encoding: utf-8 -*- |
|
3
|
|
|
# vim: set et sw=4 ts=4 sts=4 ff=unix fenc=utf8: |
|
4
|
|
|
# Author: Binux<[email protected]> |
|
5
|
|
|
# http://binux.me |
|
6
|
|
|
# Created on 2016-01-17 18:32:33 |
|
7
|
|
|
|
|
8
|
|
|
import time |
|
9
|
|
|
|
|
10
|
|
|
import elasticsearch.helpers |
|
11
|
|
|
from elasticsearch import Elasticsearch |
|
12
|
|
|
from pyspider.database.base.projectdb import ProjectDB as BaseProjectDB |
|
13
|
|
|
|
|
14
|
|
|
|
|
15
|
|
|
class ProjectDB(BaseProjectDB): |
|
16
|
|
|
__type__ = 'project' |
|
17
|
|
|
|
|
18
|
|
|
def __init__(self, hosts, index='pyspider'): |
|
19
|
|
|
self.index = index |
|
20
|
|
|
self.es = Elasticsearch(hosts=hosts) |
|
21
|
|
|
|
|
22
|
|
|
self.es.indices.create(index=self.index, ignore=400) |
|
23
|
|
|
if not self.es.indices.get_mapping(index=self.index, doc_type=self.__type__): |
|
24
|
|
|
self.es.indices.put_mapping(index=self.index, doc_type=self.__type__, body={ |
|
25
|
|
|
"_all": {"enabled": False}, |
|
26
|
|
|
"properties": { |
|
27
|
|
|
"updatetime": {"type": "double"} |
|
28
|
|
|
} |
|
29
|
|
|
}) |
|
30
|
|
|
|
|
31
|
|
|
def insert(self, name, obj={}): |
|
32
|
|
|
obj = dict(obj) |
|
33
|
|
|
obj['name'] = name |
|
34
|
|
|
obj['updatetime'] = time.time() |
|
35
|
|
|
|
|
36
|
|
|
obj.setdefault('group', '') |
|
37
|
|
|
obj.setdefault('status', 'TODO') |
|
38
|
|
|
obj.setdefault('script', '') |
|
39
|
|
|
obj.setdefault('comments', '') |
|
40
|
|
|
obj.setdefault('rate', 0) |
|
41
|
|
|
obj.setdefault('burst', 0) |
|
42
|
|
|
|
|
43
|
|
|
return self.es.index(index=self.index, doc_type=self.__type__, body=obj, id=name, |
|
44
|
|
|
refresh=True) |
|
45
|
|
|
|
|
46
|
|
|
def update(self, name, obj={}, **kwargs): |
|
47
|
|
|
kwargs.update(obj) |
|
48
|
|
|
obj = self.get(name) |
|
49
|
|
|
if obj is None: |
|
50
|
|
|
return |
|
51
|
|
|
|
|
52
|
|
|
obj.update(kwargs) |
|
53
|
|
|
obj['updatetime'] = time.time() |
|
54
|
|
|
return self.es.index(index=self.index, doc_type=self.__type__, body=obj, id=name, |
|
55
|
|
|
refresh=True) |
|
56
|
|
|
|
|
57
|
|
|
def get_all(self, fields=None): |
|
58
|
|
|
for record in elasticsearch.helpers.scan(self.es, index=self.index, doc_type=self.__type__, |
|
59
|
|
|
query={'query': {"match_all": {}}}, |
|
60
|
|
|
_source_include=fields or []): |
|
61
|
|
|
yield record['_source'] |
|
62
|
|
|
|
|
63
|
|
|
def get(self, name, fields=None): |
|
64
|
|
|
ret = self.es.get(index=self.index, doc_type=self.__type__, id=name, |
|
65
|
|
|
_source_include=fields or [], ignore=404) |
|
66
|
|
|
return ret.get('_source', None) |
|
67
|
|
|
|
|
68
|
|
|
def check_update(self, timestamp, fields=None): |
|
69
|
|
|
for record in elasticsearch.helpers.scan(self.es, index=self.index, doc_type=self.__type__, |
|
70
|
|
|
query={'query': {"range": { |
|
71
|
|
|
"updatetime": {"gte": timestamp} |
|
72
|
|
|
}}}, _source_include=fields or []): |
|
73
|
|
|
yield record['_source'] |
|
74
|
|
|
|
|
75
|
|
|
def drop(self, name): |
|
76
|
|
|
return self.es.delete(index=self.index, doc_type=self.__type__, id=name, refresh=True) |
|
77
|
|
|
|
binux /
pyspider
