level0.py

import os
from typing import List, Literal, Optional, Tuple

from .common import Result, constants, request, utils

DateTimeTuple = Tuple[str, str]


def find(
    dataset: str,
    instrument: Literal["MSC", "IFS", "MCI", "HSTDM", "CPIC"],
    data_model: str = "raw",
    data_uuid: Optional[str] = None,
    obs_group: Optional[str] = None,
    obs_id: Optional[str] = None,
    detector: Optional[str] = None,
    obs_type: Optional[str] = None,
    filter: Optional[str] = None,
    obs_date: Optional[DateTimeTuple] = None,
    create_time: Optional[DateTimeTuple] = None,
    qc_status: Optional[int] = None,
    prc_status: Optional[int] = None,
    file_name: Optional[str] = None,
    ra: Optional[int] = None,
    dec: Optional[int] = None,
    radius: Optional[float] = None,
    object: Optional[str] = None,
    page: int = 1,
    limit: int = 0,
    **extra_kwargs,
) -> Result:
    """
    根据给定的参数搜索0级数据文件记录

    Args:
        data_uuid (Optional[str], optional): 数据UUID. Defaults to None.
        data_model (str): 数据模型. Defaults is 'raw'.
        instrument (str): 设备，必需为'MSC', 'IFS', 'MCI', 'HSTDM', 'CPIC'之一.
        obs_group (Optional[str], optional): 项目ID. Defaults to None.
        obs_id (Optional[str], optional): 观测ID. Defaults to None.
        detector (Optional[str], optional): 探测器. Defaults to None.
        obs_type (Optional[str], optional): 观测类型，如'01'. Defaults to None.
        filter (Optional[str], optional): 滤光片. Defaults to None.
        obs_date (Optional[DateTimeTuple], optional): 观测时间范围. Defaults to None.
        create_time (Optional[DateTimeTuple], optional): 创建时间范围. Defaults to None.
        qc_status (Optional[int], optional): QC0状态. Defaults to None.
        prc_status (Optional[int], optional): 处理状态. Defaults to None.
        file_name (Optional[str], optional): 文件名. Defaults to None.
        ra (Optional[int], optional): 目标赤经. Defaults to None.
        dec (Optional[int], optional): 目标赤纬. Defaults to None.
        radius (Optional[float], optional): 搜索半径. Defaults to None.
        object (Optional[str], optional): 目标名称. Defaults to None.
        dataset (Optional[str], optional): 数据集名称. Defaults to constants.DEFAULT_DATASET.
        page (int, optional): 页码. Defaults to 1.
        limit (int, optional): 每页数量. Defaults to 0，不限制.

    Returns:
        Result: 搜索结果对象.

    """

    params = {
        "data_uuid": data_uuid,
        "data_model": data_model,
        "obs_group": obs_group,
        "obs_id": obs_id,
        "instrument": instrument,
        "detector": detector,
        "obs_type": obs_type,
        "filter": filter,
        "qc_status": qc_status,
        "prc_status": prc_status,
        "file_name": file_name,
        "ra_obj": ra,
        "dec_obj": dec,
        "radius": radius,
        "object": object,
        "obs_date_start": None,
        "obs_date_end": None,
        "create_time_start": None,
        "create_time_end": None,
        "dataset": dataset,
        "page": page,
        "limit": limit,
    }
    params.update(extra_kwargs)

    if obs_date is not None:
        params["obs_date_start"], params["obs_date_end"] = obs_date
        if params["obs_date_start"] and utils.is_valid_datetime_format(
            params["obs_date_start"]
        ):
            pass
        if params["obs_date_end"] and utils.is_valid_datetime_format(
            params["obs_date_end"]
        ):
            pass
    if create_time is not None:
        params["create_time_start"], params["create_time_end"] = create_time
        utils.is_valid_datetime_format(
            params["create_time_start"]
        ) or utils.is_valid_datetime_format(params["create_time_end"])
    return request.post("/api/level0", params)


def get_by_id(_id: str) -> Result:
    """
    根据内部ID获取0级数据

    Args:
        _id (str): 0级数据的内部ID

    Returns:
        Result: 查询结果

    """
    return request.get(f"/api/level0/_id/{_id}")


def get_by_uuid(_uuid: str) -> Result:
    """
    根据UUID获取0级数据

    Args:
        _uuid (str): 0级数据的UUID

    Returns:
        Result: 查询结果

    """
    return request.get(f"/api/level0/uuid/{_uuid}")


def find_by_level0_id(level0_id: str) -> Result:
    """
    通过 level0 的 ID 查询0级数据

    Args:
        level0_id (str): 0级数据的ID

    Returns:
        Result: 查询结果

    """
    return request.get(f"/api/level0/{level0_id}")


def update_qc_status(level0_id: str, qc_status: int, dataset: str) -> Result:
    """
    更新0级数据的QC状态

    Args:
        level0_id (str): 0级数据的ID
        qc_status (int): QC状态
        dataset (str): 数据集名称

    Returns:
        Result: 更新结果
    """
    return request.put(
        f"/api/level0/qc_status/{level0_id}",
        {"qc_status": qc_status, "dataset": dataset},
    )


def update_qc_status_by_ids(ids: List[str], qc_status: int) -> Result:
    """
    根据内部_id，批量更新0级数据的QC状态

    Args:
        ids (List[str]): 内部_id列表
        qc_status (int): QC状态

    Returns:
        Result: 更新结果
    """
    return request.put(
        "/api/level0/qc_status/batch/update", {"qc_status": qc_status, "ids": ids}
    )


def update_prc_status(
    level0_id: str, dag_run: str, prc_status: int, dataset: str
) -> Result:
    """
    更新0级数据的处理状态

    Args:
        level0_id (str): 0级数据的ID
        dag_run (str): DAG运行标识
        prc_status (int): 处理状态
        dataset (str): 数据集名称

    Returns:
        Result: 操作结果
    """
    return request.put(
        f"/api/level0/prc_status/{level0_id}/{dag_run}",
        {"prc_status": prc_status, "dataset": dataset},
    )


def update_prc_status_by_ids(ids: List[str], prc_status: int) -> Result:
    """
    根据内部_id，批量更新0级数据的处理状态

    Args:
        ids (List[str]): 内部_id列表
        prc_status (int): 处理状态

    Returns:
        Result: 操作结果
    """
    return request.put(
        "/api/level0/prc_status/batch", {"prc_status": prc_status, "ids": ids}
    )


def write(
    local_file: str, dataset: str = constants.DEFAULT_DATASET, **kwargs
) -> Result:
    """
    将本地文件写入DFS中

    Args:
        local_file (str]): 文件路径
        dataset (Optional[str], optional): 数据集名称. Defaults to None.
        **kwargs: 额外的关键字参数，这些参数将传递给DFS

    Returns:
        Result: 操作的结果对象，包含操作是否成功以及相关的错误信息，成功返回数据对象

    """
    params = {
        "dataset": dataset,
    }
    params.update(kwargs)
    if not os.path.exists(local_file):
        raise FileNotFoundError(local_file)
    return request.post_file("/api/level0/file", local_file, params)


def write_cat(
    local_file: str, dataset: str = constants.DEFAULT_DATASET, **kwargs
) -> Result:
    """
    主巡天仿真数据的星表本地文件写入DFS中

    Args:
        local_file (str]): 文件路径
        dataset (Optional[str], optional): 数据集名称. Defaults to None.
        **kwargs: 额外的关键字参数，这些参数将传递给DFS

    Returns:
        Result: 操作的结果对象，包含操作是否成功以及相关的错误信息，成功返回数据对象

    """
    params = {
        "dataset": dataset,
    }
    params.update(kwargs)
    if not os.path.exists(local_file):
        raise FileNotFoundError(local_file)
    return request.post_file("/api/level0/cat/file", local_file, params)


def find_process(
    dag: Optional[str] = None,
    dag_run: Optional[str] = None,
    batch_id: Optional[str] = None,
    level0_id: Optional[str] = None,
    dataset: Optional[str] = None,
    prc_module: Optional[str] = None,
    prc_status: Optional[int] = None,
    prc_date: Optional[DateTimeTuple] = None,
    page: int = 1,
    limit: int = 0,
) -> Result:
    """
    查询0级数据处理过程

    Args:
        dag (str): DAG标识
        dag_run (str): DAG运行标识
        batch_id (str): 批次ID
        level0_id (str): 0级数据的ID
        dataset (str): 数据集
        prc_module (str): 处理模块
        prc_status (int): 处理状态
        prc_date (DateTimeTuple): 处理时间范围
        page (int): 页码，默认为1
        limit (int): 每页数量 0: 不限制

    Returns:
        Result: 成功后，Result.data为数据列表，失败message为失败原因

    """
    params = {
        "dag": dag,
        "dag_run": dag_run,
        "batch_id": batch_id,
        "level0_id": level0_id,
        "dataset": dataset,
        "prc_module": prc_module,
        "prc_status": prc_status,
        "prc_date_start": None,
        "prc_date_end": None,
        "page": page,
        "limit": limit,
    }
    if prc_date is not None:
        params["prc_date_start"], params["prc_date_end"] = prc_date
        if params["prc_date_start"] and utils.is_valid_datetime_format(
            params["prc_date_start"]
        ):
            pass
        if params["prc_date_end"] and utils.is_valid_datetime_format(
            params["prc_date_end"]
        ):
            pass
    return request.post("/api/level0/process", params)


def add_process(
    level0_id: str,
    dag: str,
    dag_run: str,
    batch_id: Optional[str] = None,
    dataset: str = constants.DEFAULT_DATASET,
    prc_status: int = -1024,
    prc_date: str = utils.get_current_time(),
    prc_module: str = "",
    message: str = "",
) -> Result:
    """
    添加0级数据处理过程

    Args:
        level0_id (str): 0级数据的ID
        dag (str): DAG标识
        dag_run (str): DAG运行标识
        dataset (str): 数据集
        batch_id (str): 批次ID
        prc_date (str): 处理时间,格式为"YYYY-MM-DD HH:MM:SS"
        prc_status (int): 处理状态
        prc_module (str): 处理模块
        message (str): 处理消息

    Returns:
        Result: 成功后，Result.data为写入记录，失败message为失败原因

    """
    params = {
        "level0_id": level0_id,
        "dag": dag,
        "dag_run": dag_run,
        "dataset": dataset,
        "batch_id": batch_id,
        "prc_date": prc_date,
        "prc_status": prc_status,
        "prc_module": prc_module,
        "message": message,
    }
    utils.is_valid_datetime_format(prc_date)
    return request.post("/api/level0/prc", params)


def new(data: dict) -> Result:
    """
    新建0级数据，用于仿真数据测试

    Args:
        data (dict): 0级数据的字典表示

    Returns:
        Result: 成功后，Result.data为写入记录，失败message为失败原因

    """
    return request.post("/api/level0/new", data)