elasticsearch.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. from json import loads, dumps
  2. from requests.auth import HTTPBasicAuth
  3. base_url = 'http://localhost:9200'
  4. username = ''
  5. password = ''
  6. index = ''
  7. search_url = base_url + '/' + index + '/_search'
  8. query_type = 'match'
  9. custom_query_json = {}
  10. show_metadata = False
  11. categories = ['general']
  12. def init(engine_settings):
  13. if 'query_type' in engine_settings and engine_settings['query_type'] not in _available_query_types:
  14. raise ValueError('unsupported query type', engine_settings['query_type'])
  15. if index == '':
  16. raise ValueError('index cannot be empty')
  17. def request(query, params):
  18. if query_type not in _available_query_types:
  19. return params
  20. if username and password:
  21. params['auth'] = HTTPBasicAuth(username, password)
  22. params['url'] = search_url
  23. params['method'] = 'GET'
  24. params['data'] = dumps(_available_query_types[query_type](query))
  25. params['headers']['Content-Type'] = 'application/json'
  26. return params
  27. def _match_query(query):
  28. """
  29. The standard for full text queries.
  30. searx format: "key:value" e.g. city:berlin
  31. REF: https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-match-query.html
  32. """
  33. try:
  34. key, value = query.split(':')
  35. except:
  36. raise ValueError('query format must be "key:value"')
  37. return {"query": {"match": {key: {'query': value}}}}
  38. def _simple_query_string_query(query):
  39. """
  40. Accepts query strings, but it is less strict than query_string
  41. The field used can be specified in index.query.default_field in Elasticsearch.
  42. REF: https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-simple-query-string-query.html
  43. """
  44. return {'query': {'simple_query_string': {'query': query}}}
  45. def _term_query(query):
  46. """
  47. Accepts one term and the name of the field.
  48. searx format: "key:value" e.g. city:berlin
  49. REF: https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-term-query.html
  50. """
  51. try:
  52. key, value = query.split(':')
  53. except:
  54. raise ValueError('query format must be key:value')
  55. return {'query': {'term': {key: value}}}
  56. def _terms_query(query):
  57. """
  58. Accepts multiple terms and the name of the field.
  59. searx format: "key:value1,value2" e.g. city:berlin,paris
  60. REF: https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-terms-query.html
  61. """
  62. try:
  63. key, values = query.split(':')
  64. except:
  65. raise ValueError('query format must be key:value1,value2')
  66. return {'query': {'terms': {key: values.split(',')}}}
  67. def _custom_query(query):
  68. key, value = query.split(':')
  69. custom_query = custom_query_json
  70. for query_key, query_value in custom_query.items():
  71. if query_key == '{{KEY}}':
  72. custom_query[key] = custom_query.pop(query_key)
  73. if query_value == '{{VALUE}}':
  74. custom_query[query_key] = value
  75. return custom_query
  76. def response(resp):
  77. results = []
  78. resp_json = loads(resp.text)
  79. if 'error' in resp_json:
  80. raise Exception(resp_json['error'])
  81. for result in resp_json['hits']['hits']:
  82. r = {key: str(value) if not key.startswith('_') else value for key, value in result['_source'].items()}
  83. r['template'] = 'key-value.html'
  84. if show_metadata:
  85. r['metadata'] = {'index': result['_index'],
  86. 'id': result['_id'],
  87. 'score': result['_score']}
  88. results.append(r)
  89. return results
  90. _available_query_types = {
  91. # Full text queries
  92. # https://www.elastic.co/guide/en/elasticsearch/reference/current/full-text-queries.html
  93. 'match': _match_query,
  94. 'simple_query_string': _simple_query_string_query,
  95. # Term-level queries
  96. # https://www.elastic.co/guide/en/elasticsearch/reference/current/term-level-queries.html
  97. 'term': _term_query,
  98. 'terms': _terms_query,
  99. # Query JSON defined by the instance administrator.
  100. 'custom': _custom_query,
  101. }