inplace_fileset4.py

import argparse
import os
import sys

import omero.clients
from omero.cli import cli_login
from omero.rtypes import rstring, rlong
from omero.gateway import BlitzGateway

MANAGED_REPO_PATH = "/data/OMERO/ManagedRepository/"

def get_files_for_fileset(fs_path):
    filepaths = []
    for path, subdirs, files in os.walk(fs_path):
        for name in files:
            # print(os.path.join(path, name))
            # If we want to ignore chunks...
            if ".z" in name or ".xml" in name:
                filepaths.append(os.path.join(path, name))
    return filepaths


def create_fileset(conn, files):
    """Create a new Fileset from local files."""
    fileset = omero.model.FilesetI()
    for f in files:
        fpath, fname, fsize, sha1 = f
        orig = create_original_file(conn, fpath, fname, fsize, sha1)
        entry = omero.model.FilesetEntryI()
        entry.setClientPath(rstring(os.path.join(fpath, fname)))
        # print(dir(entry))
        entry.setOriginalFile(orig)
        fileset.addFilesetEntry(entry)
    return fileset


def create_original_file(conn, path, name, fileSize, shaHast):
    updateService = conn.getUpdateService()
    # create original file, set name, path, mimetype
    originalFile = omero.model.OriginalFileI()
    originalFile.setName(rstring(name))
    originalFile.setPath(rstring(path))
    # if mimetype:
    #     originalFile.mimetype = rstring(mimetype)
    originalFile.setSize(rlong(fileSize))
    # set sha1
    originalFile.setHash(rstring(shaHast))

    chk = omero.model.ChecksumAlgorithmI()
    chk.setValue(rstring(omero.model.enums.ChecksumAlgorithmSHA1160))
    originalFile.setHasher(chk)

    originalFile = updateService.saveAndReturnObject(originalFile, conn.SERVICE_OPTS)
    return originalFile


def getHash(localPath):
    with open(localPath, 'rb') as fileHandle:
        try:
            import hashlib
            hash_sha1 = hashlib.sha1
        except:
            import sha
            hash_sha1 = sha.new
        fileHandle.seek(0)
        h = hash_sha1()
        h.update(fileHandle.read())
        shaHast = h.hexdigest()
    return shaHast


def main(argv):
    parser = argparse.ArgumentParser()

    parser.add_argument('fileset_id', type=int, help='Fileset to replace')
    parser.add_argument('files_path', help="Full path to directory. E.g. /data/ManagedRepository/path/to/plate.zarr")
    args = parser.parse_args(argv)


    # This is the path to a symlink e.g. from Tonsil2.zarr
    # root_dir = "/data/OMERO/ManagedRepository/demo_52/Blitz-0-Ice.ThreadPool.Server-11/2023-03/31/10-49-23.116/Tonsil2.zarr"
    root_dir = args.files_path
    managed_repo_dir = root_dir.replace(MANAGED_REPO_PATH, "")


    with cli_login() as cli:
        conn = BlitzGateway(client_obj=cli._client)

        source_fileset = conn.getObject("Fileset", args.fileset_id)
        if source_fileset is None:
            print ('source Fileset id not found: %s' % args.fileset_id)
            sys.exit(1)

        # Create a new Fileset, using files in the root_dir
        file_paths = []
        prefix = source_fileset.templatePrefix
        print("templatePrefix", prefix)

        for file_path in get_files_for_fileset(root_dir):
            fpath, fname = os.path.split(file_path)
            # remove absolute path - just need relative path INSIDE ManagedRepo
            fpath = fpath.replace(MANAGED_REPO_PATH, "")
            # NB: extra slash!
            fpath = fpath + "/"
            fsize = os.path.getsize(file_path)
            fhash = getHash(file_path)
            print(fpath, fname, fsize, fhash)
            file_paths.append([fpath, fname, fsize, fhash])

        new_fileset = create_fileset(conn, file_paths)
        new_fileset.templatePrefix = omero.rtypes.rstring(prefix)

        update = conn.getUpdateService()

        new_fileset = update.saveAndReturnObject(new_fileset)
        print("Created new_fileset", new_fileset.id.val)

        for image in source_fileset.copyImages():
            img = image._obj
            img.fileset = omero.model.FilesetI(new_fileset.id.val, False)
            conn.getUpdateService().saveObject(img, conn.SERVICE_OPTS)

            # Print the HQL updates we need to update each Pixels to new Fileset file
            pid = image.getPixelsId()
            print(f"UPDATE pixels SET name = '.zattrs', path = '{managed_repo_dir}' where id = {pid};")


if __name__ == '__main__':
    main(sys.argv[1:])