data_object.py

"""Utility / convenience functions for data object IO."""
from __future__ import annotations

__copyright__ = 'Copyright (c) 2019-2026, Utrecht University'
__license__   = 'GPLv3, see LICENSE'

import binascii
import json
from typing import Dict, List, Tuple

import genquery
import irods_types

import constants
import error
import misc
import msi
import pathutil
import rule


def exists(ctx: rule.Context, path: str) -> bool:
    """Check if a data object with the given path exists.

    :param ctx:  Combined type of a callback and rei struct
    :param path: A data object path

    :returns: Boolean indicating if data object exists
    """
    coll_name, data_name = pathutil.chop(path)
    coll_name = misc.escape(coll_name)
    data_name = misc.escape(data_name)
    return len(list(genquery.Query(
               ctx, "DATA_ID",
               f"COLL_NAME = '{coll_name}' AND DATA_NAME = '{data_name}'",
               output=genquery.AS_LIST, limit=1, parser=genquery.Parser.GENQUERY2))) > 0


def get_properties(ctx: rule.Context, data_id: str, resource: str) -> Dict | None:
    """Retrieves default properties of a data object from iRODS.

    :param ctx:      Combined type of a callback and rei struct
    :param data_id:  Data ID of the data object
    :param resource: Name of resource

    :returns: Dictionary mapping each requested property to its retrieved value, or None if not found.
    """
    # Default properties available for retrieva
    properties = [
        "DATA_ID", "DATA_MODIFY_TIME", "DATA_OWNER_NAME", "DATA_SIZE",
        "COLL_ID", "DATA_RESC_HIER", "DATA_NAME", "COLL_NAME",
    ]

    # Retrieve data object with default properties
    query_fields = ", ".join(properties)
    iter = genquery.row_iterator(
        query_fields,
        "DATA_ID = '{}' AND DATA_RESC_HIER like '{}%'".format(data_id, resource),
        genquery.AS_LIST, ctx
    )

    # Return a None when no data object is found
    prop_dict = None

    for row in iter:
        prop_dict = dict(zip(properties, row))
        break

    return prop_dict


def owner(ctx: rule.Context, path: str) -> Tuple[str, str] | None:
    """Find the owner of a data object. Returns (name, zone) or None."""
    owners = list(genquery.row_iterator(
                  "DATA_OWNER_NAME, DATA_OWNER_ZONE",
                  "COLL_NAME = '%s' AND DATA_NAME = '%s'" % pathutil.chop(path),
                  genquery.AS_LIST, ctx))
    return tuple(owners[0]) if len(owners) > 0 else None


def size(ctx: rule.Context, path: str) -> int | None:
    """Get a data object's size in bytes.

    :param ctx:      Combined type of a callback and rei struct
    :param path:     Path to iRODS data object

    :returns: Data object's size or None if object is not found
    """
    coll_name, data_name = pathutil.chop(path)
    coll_name = misc.escape(coll_name)
    data_name = misc.escape(data_name)
    iter = genquery.Query(
        ctx, "DATA_SIZE, order_desc(DATA_MODIFY_TIME)",
        f"COLL_NAME = '{coll_name}' AND DATA_NAME = '{data_name}'",
        output=genquery.AS_LIST
    )

    for row in iter:
        return int(row[0])

    return None


def checksum(ctx: rule.Context, path: str) -> str | None:
    """Get a data object's checksum.

    :param ctx:  Combined type of a callback and rei struct
    :param path: Path to iRODS data object

    :returns: Data object's checksum or None if object is not found
    """
    coll_name, data_name = pathutil.chop(path)
    coll_name = misc.escape(coll_name)
    data_name = misc.escape(data_name)
    iter = genquery.Query(
        ctx, "DATA_CHECKSUM",
        f"COLL_NAME = '{coll_name}' AND DATA_NAME = '{data_name}'",
        output=genquery.AS_LIST
    )

    for row in iter:
        return row[0]

    return None


def has_replica_with_status(ctx: rule.Context, path: str, statuses: List) -> bool:
    """Check if data object has replica with specified replica statuses.

    :param ctx:      Combined type of a callback and rei struct
    :param path:     Path to iRODS data object
    :param statuses: List of replica status to check

    :returns: Boolean indicating if data object has replicas with specified replica statuses
    """
    coll_name, data_name = pathutil.chop(path)
    coll_name = misc.escape(coll_name)
    data_name = misc.escape(data_name)
    iter = genquery.row_iterator(
        "DATA_REPL_STATUS",
        f"COLL_NAME = '{coll_name}' AND DATA_NAME = '{data_name}'",
        genquery.AS_LIST, ctx
    )

    for row in iter:
        if constants.replica_status(int(row[0])) in statuses:
            return True

    return False


def write(ctx: rule.Context, path: str, data: str) -> None:
    """Write a string to an iRODS data object.

    This will overwrite the data object if it exists.

    :param ctx:  Combined type of a callback and rei struct
    :param path: Path to iRODS data object
    :param data: Data to write to data object
    """
    if exists(ctx, path):
        ret = msi.data_obj_open(ctx, 'openFlags=O_WRONLYO_TRUNC++++objPath=' + path, 0)
        handle = ret['arguments'][1]
    else:
        ret = msi.data_obj_create(ctx, path, '', 0)
        handle = ret['arguments'][2]

    msi.data_obj_write(ctx, handle, data, 0)
    msi.data_obj_close(ctx, handle, 0)


def read(ctx: rule.Context, path: str, max_size: int = constants.IIDATA_MAX_SLURP_SIZE) -> str:
    """Read an entire iRODS data object into a string."""
    sz = size(ctx, path)
    if sz is None:
        raise error.UUFileNotExistError('data_object.read: object does not exist ({})'
                                        .format(path))

    if sz > max_size:
        raise error.UUFileSizeError('data_object.read: file size limit exceeded ({} > {})'
                                    .format(sz, max_size))

    if sz == 0:
        # Don't bother reading an empty file.
        return ''

    ret = msi.data_obj_open(ctx, 'objPath=' + path, 0)
    handle = ret['arguments'][1]

    ret = msi.data_obj_read(ctx,
                            handle,
                            sz,
                            irods_types.BytesBuf())

    buf = ret['arguments'][2]

    # Convert BytesBuffer to string.
    ret_val = msi.bytes_buf_to_str(ctx, buf, "")
    output = ret_val["arguments"][1]
    msi.data_obj_close(ctx, handle, 0)

    return output


def copy(ctx: rule.Context, path_org: str, path_copy: str, force: bool = True) -> None:
    """Copy a data object.

    :param ctx:       Combined type of a callback and rei struct
    :param path_org:  Data object original path
    :param path_copy: Data object copy path
    :param force:     Applies "forceFlag"

    This may raise a error.UUError if the file does not exist, or when the user
    does not have write permission.
    """
    msi.data_obj_copy(ctx,
                      path_org,
                      path_copy,
                      'numThreads=1++++verifyChksum={}'.format('++++forceFlag=' if force else ''),
                      irods_types.BytesBuf())

    json_inp = {"logical_path": path_copy, "options": {"reference": path_org}}
    msi.touch(ctx, json.dumps(json_inp))


def remove(ctx: rule.Context, path: str, force: bool = False) -> None:
    """Delete a data object.

    :param ctx:   Combined type of a callback and rei struct
    :param path:  Data object path
    :param force: Applies "forceFlag"

    This may raise a error.UUError if the file does not exist, or when the user
    does not have write permission.
    """
    msi.data_obj_unlink(ctx,
                        'objPath={}{}'.format(path, '++++forceFlag=' if force else ''),
                        irods_types.BytesBuf())


def rename(ctx: rule.Context, path_org: str, path_target: str) -> None:
    """Rename data object from path_org to path_target.

    :param ctx:         Combined type of a callback and rei struct
    :param path_org:    Data object original path
    :param path_target: Data object new path

    This may raise a error.UUError if the file does not exist, or when the user
    does not have write permission.
    """
    # Get the modified date of the data object
    ret = msi.obj_stat(ctx, path_org, irods_types.RodsObjStat())
    output = ret['arguments'][1]
    modify_time = int(str(output.modifyTime))

    msi.data_obj_rename(ctx,
                        path_org,
                        path_target,
                        '0',
                        irods_types.BytesBuf())

    json_inp = {"logical_path": path_target, "options": {"seconds_since_epoch": modify_time}}
    msi.touch(ctx, json.dumps(json_inp))


def name_from_id(ctx: rule.Context, data_id: str) -> str | None:
    """Get data object name from data object id.

    :param ctx:     Combined type of a callback and rei struct
    :param data_id: Data object id

    :returns: Data object name
    """
    x = genquery.Query(ctx, "COLL_NAME, DATA_NAME", f"DATA_ID = '{data_id}'").first()
    if x is not None:
        return '/'.join(x)

    return None


def id_from_path(ctx: rule.Context, path: str) -> str:
    """Get data object id from data object path at its first appearance.

    :param ctx:  Combined type of a callback and rei struct
    :param path: Path to iRODS data object

    :returns: Data object id
    """
    coll_name, data_name = pathutil.chop(path)
    coll_name = misc.escape(coll_name)
    data_name = misc.escape(data_name)

    return genquery.Query(ctx, "DATA_ID",
                          f"COLL_NAME = '{coll_name}' AND DATA_NAME = '{data_name}'").first()


def decode_checksum(checksum: str) -> str:
    """Decode data object checksum.

    :param checksum: Base64 encoded SHA256 checksum

    :returns: Data object's SHA256 checksum
    """
    if checksum is None:
        return "0"
    else:
        return binascii.hexlify(binascii.a2b_base64(checksum[5:])).decode("UTF-8")


def get_group_owners(ctx: rule.Context, path: str) -> List:
    """Return list of groups of data object, each entry being name of the group and the zone."""
    coll_name, data_name = pathutil.chop(path)
    coll_name = misc.escape(coll_name)
    data_name = misc.escape(data_name)

    groups = list(genquery.Query(
        ctx, "USER_NAME, USER_ZONE",
        f"COLL_NAME = '{coll_name}' and DATA_NAME = '{data_name}' AND USER_TYPE = 'rodsgroup' AND DATA_ACCESS_NAME = 'own'",
        output=genquery.AS_LIST))

    return groups