Text2Query: Add RESTful API for database health check (opea-project#1935)

* Text2Query: Add RESTful API for database health check

Signed-off-by: Yi Yao <[email protected]>

---------

Signed-off-by: Yi Yao <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: joshuayao

comps/text2query/src/integrations/text2cypher.py

@@ -215,3 +215,16 @@ async def invoke(self, input: Text2QueryRequest):
             raise
 
         return result
+
+    async def db_connection_check(self, request: Text2QueryRequest):
+        """Check the connection to Neo4j database.
+
+        This function takes a Text2QueryRequest object containing the database connection information.
+        It attempts to connect to the database using the provided connection URL and credentials.
+
+        Args:
+            request (Text2QueryRequest): A Text2QueryRequest object with the database connection information.
+        Returns:
+            dict: A dictionary with a 'status' key indicating whether the connection was successful or failed.
+        """
+        return {"status": "Connection successful"}

comps/text2query/src/integrations/text2sql.py

@@ -6,9 +6,12 @@
 import os
 from urllib.parse import urlparse, urlunparse
 
+from fastapi.exceptions import HTTPException
 from langchain.agents.agent_types import AgentType
 from langchain_community.utilities.sql_database import SQLDatabase
 from langchain_huggingface import HuggingFaceEndpoint
+from sqlalchemy import create_engine
+from sqlalchemy.exc import SQLAlchemyError
 
 from comps import CustomLogger, OpeaComponent, OpeaComponentRegistry, ServiceType
 from comps.cores.proto.api_protocol import Text2QueryRequest
@@ -102,7 +105,9 @@ def format_db_url(self, request: Text2QueryRequest) -> str:
 
     async def invoke(self, request: Text2QueryRequest):
         url = request.conn_url
-        if not url:
+        if url:
+            url = self.format_db_url(request)
+        else:
             raise ValueError("Database connection URL must be provided in 'conn_url' field of the request.")
 
         """Execute a SQL query using the custom SQL agent.
@@ -132,3 +137,29 @@ async def invoke(self, request: Text2QueryRequest):
                 query.append(log.tool_input)
         result["sql"] = query[0].replace("Observation", "")
         return {"result": result}
+
+    async def db_connection_check(self, request: Text2QueryRequest):
+        """Check the connection to the database.
+
+        This function takes a Text2QueryRequest object containing the database connection information.
+        It attempts to connect to the database using the provided connection URL and credentials.
+
+        Args:
+            request (Text2QueryRequest): A Text2QueryRequest object with the database connection information.
+        Returns:
+            dict: A dictionary with a 'status' key indicating whether the connection was successful or failed.
+        """
+        url = request.conn_url
+        if url:
+            url = self.format_db_url(request)
+        else:
+            raise ValueError("Database connection URL must be provided in 'conn_url' field of the request.")
+
+        try:
+            engine = create_engine(url)
+            with engine.connect() as _:
+                # If the connection is successful, return True
+                return {"status": "Connection successful"}
+        except SQLAlchemyError as e:
+            logger.error(f"Connection failed: {e}")
+            raise HTTPException(status_code=500, detail=f"Failed to connect to database: {url}")

comps/text2query/src/opea_text2query_loader.py

@@ -0,0 +1,10 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from comps import OpeaComponentLoader
+
+
+class OpeaText2QueryLoader(OpeaComponentLoader):
+
+    async def db_connection_check(self, *args, **kwargs):
+        return await self.component.db_connection_check(*args, **kwargs)

comps/text2query/src/opea_text2query_microservice.py

@@ -13,8 +13,12 @@
 
 import os
 
-from comps import CustomLogger, OpeaComponentLoader, opea_microservices, register_microservice
+from fastapi import status
+from fastapi.exceptions import HTTPException
+
+from comps import CustomLogger, opea_microservices, register_microservice
 from comps.cores.proto.api_protocol import Text2QueryRequest
+from comps.text2query.src.opea_text2query_loader import OpeaText2QueryLoader
 
 logger = CustomLogger("text2query")
 logflag = os.getenv("LOGFLAG", False)
@@ -34,7 +38,7 @@
     raise ValueError(f"Unsupported TEXT2QUERY_COMPONENT_NAME: {component_name}")
 
 # Initialize the OPEA component loader with the selected component
-loader = OpeaComponentLoader(
+loader = OpeaText2QueryLoader(
     component_name,
     description=f"OPEA TEXT2QUERY Component: {component_name}",
 )
@@ -63,6 +67,32 @@ async def execute_agent(request: Text2QueryRequest):
     return await loader.invoke(request)
 
 
+@register_microservice(
+    name="opea_service@text2query",
+    endpoint="/v1/db/health",
+    host="0.0.0.0",
+    port=9097,
+)
+async def db_connection_check(request: Text2QueryRequest):
+    """Check the connection to the database.
+
+    This function takes an Input object containing the database connection information.
+    It uses the test_connection method of the PostgresConnection class to check if the connection is successful.
+
+    Args:
+        request (Text2QueryRequest): An Input object with the database connection information.
+
+    Returns:
+        dict: A dictionary with a 'status' key indicating whether the connection was successful or failed.
+    """
+    logger.info(f"Received input for connection check: {request}")
+    if not isinstance(request, Text2QueryRequest):
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Input type mismatch: expected Text2QueryRequest"
+        )
+    return await loader.db_connection_check(request)
+
+
 if __name__ == "__main__":
     logger.info("OPEA Text2Query Microservice is starting...")
     opea_microservices["opea_service@text2query"].start()

tests/text2query/test_text2query_text2sql.sh

@@ -69,6 +69,23 @@ function start_service() {
 
 function validate_microservice() {
     url="postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${ip_address}:5442/${POSTGRES_DB}"
+
+    echo "Validating v1/db/health..."
+    result=$(http_proxy="" curl http://${ip_address}:${TEXT2SQL_PORT}/v1/db/health\
+        -X POST \
+        -d '{"conn_type": "sql", "conn_url": "'${url}'", "conn_user": "'${POSTGRES_USER}'","conn_password": "'${POSTGRES_PASSWORD}'","conn_dialect": "postgresql" }' \
+        -H 'Content-Type: application/json')
+
+    if [[ $result == *"Connection successful"* ]]; then
+        echo $result
+        echo "Result correct."
+    else
+        echo "Result wrong. Received was $result"
+        docker logs text2query-sql-server > ${LOG_PATH}/text2query.log
+        exit 1
+    fi
+
+    echo "Validating v1/text2query..."
     result=$(http_proxy="" curl http://${ip_address}:${TEXT2SQL_PORT}/v1/text2query\
         -X POST \
         -d '{"query": "Find the total number of Albums.","conn_type": "sql", "conn_url": "'${url}'", "conn_user": "'${POSTGRES_USER}'","conn_password": "'${POSTGRES_PASSWORD}'","conn_dialect": "postgresql" }' \
@@ -83,7 +100,6 @@ function validate_microservice() {
         docker logs tgi-server > ${LOG_PATH}/tgi.log
         exit 1
     fi
-
 }
 
 function stop_docker() {

tests/text2query/test_text2query_text2sql_on_intel_hpu.sh

@@ -68,6 +68,23 @@ function start_service() {
 
 function validate_microservice() {
     url="postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${ip_address}:5442/${POSTGRES_DB}"
+
+    echo "Validating v1/db/health..."
+    result=$(http_proxy="" curl http://${ip_address}:${TEXT2SQL_PORT}/v1/db/health\
+        -X POST \
+        -d '{"conn_type": "sql", "conn_url": "'${url}'", "conn_user": "'${POSTGRES_USER}'","conn_password": "'${POSTGRES_PASSWORD}'","conn_dialect": "postgresql" }' \
+        -H 'Content-Type: application/json')
+
+    if [[ $result == *"Connection successful"* ]]; then
+        echo $result
+        echo "Result correct."
+    else
+        echo "Result wrong. Received was $result"
+        docker logs text2query-sql-server > ${LOG_PATH}/text2query.log
+        exit 1
+    fi
+
+    echo "Validating v1/text2query..."
     result=$(http_proxy="" curl http://${ip_address}:${TEXT2SQL_PORT}/v1/text2query\
         -X POST \
         -d '{"query": "Find the total number of Albums.","conn_type": "sql", "conn_url": "'${url}'", "conn_user": "'${POSTGRES_USER}'","conn_password": "'${POSTGRES_PASSWORD}'","conn_dialect": "postgresql" }' \