-
Notifications
You must be signed in to change notification settings - Fork 0
/
es.py
76 lines (63 loc) · 1.7 KB
/
es.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import os
import boto3
from elasticsearch import Elasticsearch, RequestsHttpConnection
from requests_aws4auth import AWS4Auth
host = os.getenv("ES_HOST", "localhost")
if host == "localhost":
ES = Elasticsearch()
else:
credentials = boto3.Session().get_credentials()
awsauth = AWS4Auth(
credentials.access_key,
credentials.secret_key,
os.getenv("AWS_DEFAULT_REGION", "us-east-1"),
"es",
session_token=credentials.token,
)
ES = Elasticsearch(
hosts=[{"host": host, "port": 443}],
http_auth=awsauth,
use_ssl=True,
verify_certs=True,
connection_class=RequestsHttpConnection,
)
ES.index(
index="test-index",
body={"__init__": True},
refresh=True,
)
def get_doc(query):
res = ES.search(
index="test-index",
body={"query": {"query_string": {"query": f'"{query}"'}}},
)
source = res["hits"]["hits"][0]["_source"] if res["hits"]["hits"] else {}
if source:
return source, res["hits"]["hits"][0]["_id"]
res = ES.search(
index="test-index",
body={"query": {"match": {"key": {"query": query}}}},
)
source = res["hits"]["hits"][0]["_source"] if res["hits"]["hits"] else {}
id_ = res["hits"]["hits"][0]["_id"] if res["hits"]["hits"] else None
return source, id_
def update_doc(source, id_):
ES.update(
index="test-index",
doc_type="_doc",
id=id_,
body={"doc": source},
refresh=True,
)
def add_doc(body):
ES.index(
index="test-index",
body=body,
refresh=True,
)
def delete_doc(id_):
ES.delete(
index="test-index",
id=id_,
refresh=True,
)