add db auth

Author: yaozheng-fang

veadk/auth/veauth/opensearch_veauth.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from typing_extensions import override
+
+from veadk.auth.veauth.base_veauth import BaseVeAuth
+from veadk.utils.logger import get_logger
+
+# from veadk.utils.volcengine_sign import ve_request
+
+logger = get_logger(__name__)
+
+
+class OpensearchVeAuth(BaseVeAuth):
+    def __init__(
+        self,
+        access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
+        secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
+    ) -> None:
+        super().__init__(access_key, secret_key)
+
+        self._token: str = ""
+
+    @override
+    def _fetch_token(self) -> None:
+        logger.info("Fetching Opensearch STS token...")
+
+        # res = ve_request(
+        #     request_body={},
+        #     action="GetOrCreatePromptPilotAPIKeys",
+        #     ak=self.access_key,
+        #     sk=self.secret_key,
+        #     service="ark",
+        #     version="2024-01-01",
+        #     region="cn-beijing",
+        #     host="open.volcengineapi.com",
+        # )
+        # try:
+        #     self._token = res["Result"]["APIKeys"][0]["APIKey"]
+        # except KeyError:
+        #     raise ValueError(f"Failed to get Prompt Pilot token: {res}")
+
+    @property
+    def token(self) -> str:
+        if self._token:
+            return self._token
+        self._fetch_token()
+        return self._token

veadk/auth/veauth/postgresql_veauth.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from typing_extensions import override
+
+from veadk.auth.veauth.base_veauth import BaseVeAuth
+from veadk.utils.logger import get_logger
+
+# from veadk.utils.volcengine_sign import ve_request
+
+logger = get_logger(__name__)
+
+
+class PostgreSqlVeAuth(BaseVeAuth):
+    def __init__(
+        self,
+        access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
+        secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
+    ) -> None:
+        super().__init__(access_key, secret_key)
+
+        self._token: str = ""
+
+    @override
+    def _fetch_token(self) -> None:
+        logger.info("Fetching PostgreSQL STS token...")
+
+        # res = ve_request(
+        #     request_body={},
+        #     action="GetOrCreatePromptPilotAPIKeys",
+        #     ak=self.access_key,
+        #     sk=self.secret_key,
+        #     service="ark",
+        #     version="2024-01-01",
+        #     region="cn-beijing",
+        #     host="open.volcengineapi.com",
+        # )
+        # try:
+        #     self._token = res["Result"]["APIKeys"][0]["APIKey"]
+        # except KeyError:
+        #     raise ValueError(f"Failed to get Prompt Pilot token: {res}")
+
+    @property
+    def token(self) -> str:
+        if self._token:
+            return self._token
+        self._fetch_token()
+        return self._token

veadk/configs/database_configs.py

@@ -32,6 +32,8 @@ class OpensearchConfig(BaseSettings):
 
     password: str = ""
 
+    secret_token: str = ""
+
 
 class MysqlConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_MYSQL_")
@@ -46,6 +48,9 @@ class MysqlConfig(BaseSettings):
 
     charset: str = "utf8"
 
+    secret_token: str = ""
+    """STS token for MySQL auth, not supported yet."""
+
 
 class PostgreSqlConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_POSTGRESQL_")
@@ -60,6 +65,8 @@ class PostgreSqlConfig(BaseSettings):
 
     database: str = ""
 
+    secret_token: str = ""
+
 
 class RedisConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_REDIS_")
@@ -72,6 +79,9 @@ class RedisConfig(BaseSettings):
 
     db: int = 0
 
+    secret_token: str = ""
+    """STS token for Redis auth, not supported yet."""
+
 
 class VikingKnowledgebaseConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_VIKING_")

veadk/memory/short_term_memory.py

@@ -30,8 +30,6 @@
 )
 from veadk.utils.logger import get_logger
 
-# from .short_term_memory_processor import ShortTermMemoryProcessor
-
 logger = get_logger(__name__)
 
 DEFAULT_LOCAL_DATABASE_PATH = "/tmp/veadk_local_database.db"
@@ -71,6 +69,7 @@ def model_post_init(self, __context: Any) -> None:
                 logger.warning(
                     "Backend `database` is deprecated, use `sqlite` to create short term memory."
                 )
+                self.backend = "sqlite"
             match self.backend:
                 case "local":
                     self.session_service = InMemorySessionService()
@@ -101,7 +100,7 @@ async def create_session(
         app_name: str,
         user_id: str,
         session_id: str,
-    ):
+    ) -> None:
         if isinstance(self.session_service, DatabaseSessionService):
             list_sessions_response = await self.session_service.list_sessions(
                 app_name=app_name, user_id=user_id