添加新的数据源OpenSearch v0.5-beta (#2759)

* 添加favicon图片

* firset

* 修改

* 撤销

* add open search

* add open search

* add open

* elasticsearch

* 修改路径错误

* 修改路径错误

* 修改路径错误 撤回

* 修改路径错误 撤回

* 添加测试方法

* 格式化

* 单元测试方法

* 重新提交

---------

Co-authored-by: 王飞 <fei.wang@xgo.one>

archery/settings.py

@@ -54,6 +54,7 @@
             "cassandra",
             "doris",
             "elasticsearch",
+            "opensearch",
         ],
     ),
     ENABLED_NOTIFIERS=(
@@ -103,6 +104,7 @@
     "odps": {"path": "sql.engines.odps:ODPSEngine"},
     "doris": {"path": "sql.engines.doris:DorisEngine"},
     "elasticsearch": {"path": "sql.engines.elasticsearch:ElasticsearchEngine"},
+    "opensearch": {"path": "sql.engines.elasticsearch:OpenSearchEngine"},
 }
 
 ENABLED_NOTIFIERS = env("ENABLED_NOTIFIERS")

requirements.txt

@@ -45,4 +45,5 @@ cassandra-driver
 httpx
 OpenAI
 elasticsearch==8.14.0
+opensearch_py==2.6.0
 

sql/engines/elasticsearch.py

@@ -2,6 +2,7 @@
 import logging
 import re
 import traceback
+from opensearchpy import OpenSearch
 import simplejson as json
 from . import EngineBase
 from .models import ResultSet, ReviewSet, ReviewResult
@@ -407,3 +408,44 @@ def get_connection(self, db_name=None):
         if not self.conn:
             raise Exception("Elasticsearch 连接无法建立。")
         return self.conn
+
+
+class OpenSearchEngine(ElasticsearchEngineBase):
+    """OpenSearch 引擎实现"""
+
+    def __init__(self, instance=None):
+        super().__init__(instance=instance)
+
+    name: str = "OpenSearch"
+    info: str = "OpenSearch 引擎"
+
+    def get_connection(self, db_name=None):
+        if self.conn:
+            return self.conn
+        if self.instance:
+            scheme = "https" if self.is_ssl else "http"
+            hosts = [
+                {
+                    "host": self.host,
+                    "port": self.port,
+                    "scheme": scheme,
+                    "use_ssl": self.is_ssl,
+                }
+            ]
+            http_auth = (
+                (self.user, self.password) if self.user and self.password else None
+            )
+            self.db_name = (self.db_name or "") + "*"
+
+            try:
+                # 创建 OpenSearch 连接
+                self.conn = OpenSearch(
+                    hosts=hosts,
+                    http_auth=http_auth,
+                    verify_certs=True,  # 开启证书验证
+                )
+            except Exception as e:
+                raise Exception(f"OpenSearch 连接建立失败: {str(e)}")
+        if not self.conn:
+            raise Exception("OpenSearch 连接无法建立。")
+        return self.conn

sql/engines/test_opensearch.py

@@ -0,0 +1,41 @@
+import json
+import unittest
+from unittest.mock import patch, Mock
+from sql.engines import ResultSet, ReviewSet
+from sql.engines.elasticsearch import OpenSearchEngine
+from sql.models import Instance
+
+
+class TestOpenSearchEngine(unittest.TestCase):
+    def setUp(self):
+        # 创建一个模拟的 instance 对象，包含必要的属性
+        self.mock_instance = Instance()
+        self.mock_instance.host = "localhost"
+        self.mock_instance.port = 9200
+        self.mock_instance.user = "user"
+        self.mock_instance.password = "pass"
+        self.mock_instance.is_ssl = True
+
+        # 初始化 OpenSearchEngine
+        self.engine = OpenSearchEngine(instance=self.mock_instance)
+
+    @patch("sql.engines.elasticsearch.OpenSearch")
+    def test_get_all_databases(self, mockElasticsearch):
+        mock_conn = Mock()
+        mock_conn.indices.get_alias.return_value = {
+            "test__index1": {},
+            "test__index2": {},
+            ".kibana_2": {},
+            ".internal.index": {},
+        }
+        mockElasticsearch.return_value = mock_conn
+
+        result = self.engine.get_all_databases()
+        expected_result = [
+            "other",
+            "system",
+            "system_internal",
+            "system_kibana",
+            "test",
+        ]
+        self.assertEqual(result.rows, expected_result)

sql/models.py

@@ -135,6 +135,7 @@ class Meta:
     ("cassandra", "Cassandra"),
     ("doris", "Doris"),
     ("elasticsearch", "Elasticsearch"),
+    ("opensearch", "OpenSearch"),
 )