DB-GPT/dbgpt/client/datasource.py

"""this module contains the datasource client functions."""
from typing import List

from dbgpt._private.pydantic import model_to_dict
from dbgpt.core.schema.api import Result

from .client import Client, ClientException
from .schema import DatasourceModel


async def create_datasource(
    client: Client, datasource: DatasourceModel
) -> DatasourceModel:
    """Create a new datasource.

    Args:
        client (Client): The dbgpt client.
        datasource (DatasourceModel): The datasource model.
    """
    try:
        res = await client.get("/datasources", model_to_dict(datasource))
        result: Result = res.json()
        if result["success"]:
            return DatasourceModel(**result["data"])
        else:
            raise ClientException(status=result["err_code"], reason=result)
    except Exception as e:
        raise ClientException(f"Failed to create datasource: {e}")


async def update_datasource(
    client: Client, datasource: DatasourceModel
) -> DatasourceModel:
    """Update a datasource.

    Args:
        client (Client): The dbgpt client.
        datasource (DatasourceModel): The datasource model.
    Returns:
        DatasourceModel: The datasource model.
    Raises:
        ClientException: If the request failed.
    """
    try:
        res = await client.put("/datasources", model_to_dict(datasource))
        result: Result = res.json()
        if result["success"]:
            return DatasourceModel(**result["data"])
        else:
            raise ClientException(status=result["err_code"], reason=result)
    except Exception as e:
        raise ClientException(f"Failed to update datasource: {e}")


async def delete_datasource(client: Client, datasource_id: str) -> DatasourceModel:
    """
    Delete a datasource.

    Args:
        client (Client): The dbgpt client.
        datasource_id (str): The datasource id.
    Returns:
        DatasourceModel: The datasource model.
    Raises:
        ClientException: If the request failed.
    """
    try:
        res = await client.delete("/datasources/" + datasource_id)
        result: Result = res.json()
        if result["success"]:
            return DatasourceModel(**result["data"])
        else:
            raise ClientException(status=result["err_code"], reason=result)
    except Exception as e:
        raise ClientException(f"Failed to delete datasource: {e}")


async def get_datasource(client: Client, datasource_id: str) -> DatasourceModel:
    """
    Get a datasource.

    Args:
        client (Client): The dbgpt client.
        datasource_id (str): The datasource id.
    Returns:
        DatasourceModel: The datasource model.
    Raises:
        ClientException: If the request failed.
    """
    try:
        res = await client.get("/datasources/" + datasource_id)
        result: Result = res.json()
        if result["success"]:
            return DatasourceModel(**result["data"])
        else:
            raise ClientException(status=result["err_code"], reason=result)
    except Exception as e:
        raise ClientException(f"Failed to get datasource: {e}")


async def list_datasource(client: Client) -> List[DatasourceModel]:
    """
    List datasources.

    Args:
        client (Client): The dbgpt client.
    Returns:
        List[DatasourceModel]: The list of datasource models.
    Raises:
        ClientException: If the request failed.
    """
    try:
        res = await client.get("/datasources")
        result: Result = res.json()
        if result["success"]:
            return [DatasourceModel(**datasource) for datasource in result["data"]]
        else:
            raise ClientException(status=result["err_code"], reason=result)
    except Exception as e:
        raise ClientException(f"Failed to list datasource: {e}")