mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
# Author: Junjun
|
|
# Date: 2025/9/9
|
|
|
|
import json
|
|
|
|
import requests
|
|
from elasticsearch import Elasticsearch
|
|
|
|
from apps.datasource.models.datasource import DatasourceConf
|
|
|
|
|
|
def get_es_connect(conf: DatasourceConf):
|
|
es_client = Elasticsearch(
|
|
[conf.host], # ES address
|
|
basic_auth=(conf.username, conf.password),
|
|
verify_certs=False,
|
|
compatibility_mode=True
|
|
)
|
|
return es_client
|
|
|
|
|
|
# get tables
|
|
def get_es_index(conf: DatasourceConf):
|
|
es_client = get_es_connect(conf)
|
|
indices = es_client.cat.indices(format="json")
|
|
res = []
|
|
if indices is not None:
|
|
for idx in indices:
|
|
index_name = idx.get('index')
|
|
desc = ''
|
|
# get mapping
|
|
mapping = es_client.indices.get_mapping(index=index_name)
|
|
mappings = mapping.get(index_name).get("mappings")
|
|
if mappings.get('_meta'):
|
|
desc = mappings.get('_meta').get('description')
|
|
res.append((index_name, desc))
|
|
return res
|
|
|
|
|
|
# get fields
|
|
def get_es_fields(conf: DatasourceConf, table_name: str):
|
|
es_client = get_es_connect(conf)
|
|
index_name = table_name
|
|
mapping = es_client.indices.get_mapping(index=index_name)
|
|
properties = mapping.get(index_name).get("mappings").get("properties")
|
|
res = []
|
|
if properties is not None:
|
|
for field, config in properties.items():
|
|
field_type = config.get("type")
|
|
desc = ''
|
|
if config.get("_meta"):
|
|
desc = config.get("_meta").get('description')
|
|
|
|
if field_type:
|
|
res.append((field, field_type, desc))
|
|
else:
|
|
# object、nested...
|
|
res.append((field, ','.join(list(config.keys())), desc))
|
|
return res
|
|
|
|
|
|
def get_es_data(conf: DatasourceConf, sql: str, table_name: str):
|
|
r = requests.post(f"{conf.host}/_sql/translate", json={"query": sql})
|
|
# print(json.dumps(r.json()))
|
|
|
|
es_client = get_es_connect(conf)
|
|
response = es_client.search(
|
|
index=table_name,
|
|
body=json.dumps(r.json())
|
|
)
|
|
|
|
# print(response)
|
|
fields = get_es_fields(conf, table_name)
|
|
res = []
|
|
for hit in response.get('hits').get('hits'):
|
|
item = []
|
|
if 'fields' in hit:
|
|
result = hit.get('fields') # {'title': ['Python'], 'age': [30]}
|
|
for field in fields:
|
|
v = result.get(field[0])
|
|
item.append(v[0]) if v else item.append(None)
|
|
res.append(tuple(item))
|
|
# print(hit['fields']['title'][0])
|
|
# elif '_source' in hit:
|
|
# print(hit.get('_source'))
|
|
return res, fields
|