def test__generate_location_query_item__only_table(self):
elastic_index_connector = ElasticSingleIndexConnector(schema_with_table)
basic_query = BasicQuery([], table_name="test")
location_basic_query = elastic_index_connector.generate_location_query_item(basic_query)
location_string_query = elastic_index_connector.generate_query_item(location_basic_query)
table_query_item = QueryItem("_table", "test")
result_string_query = elastic_index_connector.generate_query_item(BasicQuery([table_query_item]))
self.assertEqual(location_string_query, result_string_query)
def test__generate_basic_query_location__basic_location(self):
elastic_index_connector = ElasticMultiIndexConnector(empty_schema)
basic_query = BasicQuery([], table_name="test")
location_query = elastic_index_connector.generate_basic_query_location(
basic_query)
self.assertDictEqual(location_query, {
"index": "test",
})
def test__generate_basic_query__multi_query(self):
elastic_index_connector = ElasticBaseConnector(empty_schema)
item = [QueryItem("test", "value"), QueryItem("test_two", "test_two")]
basic_query = BasicQuery(item, operator="AND")
generate_query = elastic_index_connector.generate_basic_query(basic_query)
result = elastic_template_query
result["query"]["query_string"]["query"] = "((test=value) AND (test_two=test_two))"
self.assertDictEqual(generate_query, result)
def test__generate_basic_query__with_fields(self):
elastic_index_connector = ElasticBaseConnector(empty_schema)
item = [QueryItem("test", "value")]
return_list = ["test1", "test2"]
basic_query = BasicQuery(item, return_list=return_list, operator="AND")
generate_query = elastic_index_connector.generate_basic_query(basic_query)
result = elastic_template_query
result["query"]["query_string"]["query"] = "(test=value)"
result["_source"] = return_list
self.assertDictEqual(generate_query, result)
def generate_location_query_item(self, basic_query):
location_query = []
if self.schema.get("table_key"):
table_query = QueryItem(self.schema.get("table_key"),
basic_query.table_name,
{FIXED_FIELD: True})
location_query.append(table_query)
if self.schema.get("collection_key"):
collection_query = QueryItem(self.schema.get("collection_key"),
basic_query.collection_name,
{"fixed": True})
location_query.append(collection_query)
return BasicQuery(location_query)
def test__get_query_body__one_query(self):
elastic_index_connector = ElasticSingleIndexConnector(schema_with_table)
return_list = ["test1", "test2"]
item = [QueryItem("test", "value"), QueryItem("test_two", "test_two")]
basic_query = BasicQuery(item, return_list=return_list, operator="AND", table_name="test")
advance_query = AdvanceQuery(basic_query)
generate_query = elastic_index_connector.get_query_body(advance_query)
data_query = elastic_template_query
data_query["query"]["query_string"]["query"] = "((_table=test) AND (test=value) AND (test_two=test_two))"
data_query["_source"] = return_list
location = {
"index": "main_index",
}
self.assertListEqual(generate_query, [location, data_query])
def get_query_body(self, advance_query: AdvanceQuery) -> List[dict]:
body = []
for basic_query in advance_query.basic_queries:
body.append(self.generate_basic_query_location(basic_query))
body.append(self.generate_basic_query(basic_query))
return body
async def query_data(self, advance_query: AdvanceQuery):
elastic_response = await self.es.msearch(
body=self.get_query_body(advance_query))
result_hits = map(self.extract_row,
elastic_response.get("responses", []))
all_results = reduce(lambda prev, value: prev + value, result_hits)
return map(self.parse_result, all_results)
if __name__ == '__main__':
empty_schema = {
"connection": {
"hosts": ["http://localhost:9200"],
"elasticOption": {
"http_auth": ("elastic", "changeme")
}
}
}
elastic_index_connector = ElasticMultiIndexConnector(empty_schema)
query_items = [QueryItem("id", "24")]
basic_queries = BasicQuery(query_items, table_name="spotify-data")
advance_query = AdvanceQuery(basic_queries)
def test__generate_query_item__basic_query(self):
elastic_index_connector = ElasticBaseConnector(empty_schema)
item = QueryItem("test", "value")
basic_query = BasicQuery([item])
generate_query = elastic_index_connector.generate_query_item(basic_query)
self.assertEqual(generate_query, "(test=value)")