From e1fc3f61d14b55fc6d9c574702069745f7dded06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E6=99=93=E9=A3=9E?= <feiazifeiazi@hotmail.com>
Date: Thu, 15 Aug 2024 15:19:51 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=96=B0=E7=9A=84=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E6=BA=90OpenSearch=20v0.5-beta=20(#2759)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 添加favicon图片

* firset

* 修改

* 撤销

* add open search

* add open search

* add open

* elasticsearch

* 修改路径错误

* 修改路径错误

* 修改路径错误 撤回

* 修改路径错误 撤回

* 添加测试方法

* 格式化

* 单元测试方法

* 重新提交

---------

Co-authored-by: 王飞 <fei.wang@xgo.one>
---
 archery/settings.py            |  2 ++
 requirements.txt               |  1 +
 sql/engines/elasticsearch.py   | 42 ++++++++++++++++++++++++++++++++++
 sql/engines/test_opensearch.py | 41 +++++++++++++++++++++++++++++++++
 sql/models.py                  |  1 +
 5 files changed, 87 insertions(+)
 create mode 100644 sql/engines/test_opensearch.py

diff --git a/archery/settings.py b/archery/settings.py
index 949b0640fe..eca54265d5 100644
--- a/archery/settings.py
+++ b/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")
diff --git a/requirements.txt b/requirements.txt
index 9555e008b1..0a2c8bef2e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -45,4 +45,5 @@ cassandra-driver
 httpx
 OpenAI
 elasticsearch==8.14.0
+opensearch_py==2.6.0
 
diff --git a/sql/engines/elasticsearch.py b/sql/engines/elasticsearch.py
index 7a08cf40f5..99eb09c301 100644
--- a/sql/engines/elasticsearch.py
+++ b/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
diff --git a/sql/engines/test_opensearch.py b/sql/engines/test_opensearch.py
new file mode 100644
index 0000000000..a67da09c1a
--- /dev/null
+++ b/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)
diff --git a/sql/models.py b/sql/models.py
index f72cbdc845..460a43e34f 100755
--- a/sql/models.py
+++ b/sql/models.py
@@ -135,6 +135,7 @@ class Meta:
     ("cassandra", "Cassandra"),
     ("doris", "Doris"),
     ("elasticsearch", "Elasticsearch"),
+    ("opensearch", "OpenSearch"),
 )