commit

datatrail.commit.commit_api

Implement the following REST endpoints.

1. *Create a dataset.
2. *Register files to a dataset.
3. *Update transformation status of file replicas.
4. Fetch active replications.
5. Fetch upcoming replications.
6. Fetch datasets.
7. Fetch content of a dataset.

add_file_replicas(session, payload, test_mode=False) async

Add file replicas.

Source code in datatrail/commit/commit_api.py

async def add_file_replicas(
    session: Union[AsyncSession, sessionmaker],
    payload: Dict[str, Any],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Add file replicas."""
    try:
        storage_map = db_queries.get_storage_map(session)
        # TODO Check if files already exist for this dataset.
        replicas = []
        for fr in payload["file_replicas"]:
            r = common.FileReplica(**fr).dict()
            r["deletion_state"] = models.TransformationState.available
            r["deletion_state_updated_at"] = r["date_created"]
            r["replication_state"] = models.TransformationState.available
            r["replication_state_updated_at"] = r["date_created"]
            replicas.append(r)
        status = db_queries.bulk_insert_file_replicas(
            session,
            replicas,
            payload["storage_name"],
            storage_map[payload["storage_element_captured_at"]],
        )

        # identify number of replicas for each file to increament
        replica_map: Dict[str, int] = {}
        for fr in replicas:
            if fr["file_name"] in replica_map:
                replica_map[fr["file_name"]] += 1
            else:
                replica_map[fr["file_name"]] = 1
        names = [(file_name,) for file_name in replica_map]
        files = db_queries.fetch_files(
            session, names, storage_map[payload["storage_element_captured_at"]]
        )
        for f in files:
            f.num_available_replicas = len(
                [
                    fr
                    for fr in f.file_replicas
                    if fr.deletion_state != models.TransformationState.completed
                ]
            )
        session.commit()

        # Now apply the policy to the file replicas
        ds = db_queries.fetch_dataset(
            session, payload["dataset_name"], payload["dataset_scope"]
        )
        rp, dp = policy_evaluator.determine_policies(session, ds, storage_map)
        for f in files:
            file_replicas = [i for i in f.file_replicas]
            replica_locations = policy_evaluator.get_replica_locations(
                session, f, storage_map
            )
            if len(replicas):
                policy_evaluator.apply_policy_to_dataset_file_replicas(
                    file_replicas,
                    replica_locations,
                    deepcopy(rp),
                    deepcopy(dp),
                    storage_map,
                )
        session.commit()
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

attach_datasets(session, payload, test_mode=False) async

Attach two datasets.

Source code in datatrail/commit/commit_api.py

async def attach_datasets(
    session: Union[AsyncSession, sessionmaker],
    payload: Union[common.DeletionPayload, common.ReplicationPayload],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Attach two datasets."""
    try:
        dis = [datasets.DatasetIdentifier(**bt) for bt in payload["belongs_to"]]
        ad = datasets.AttachDataset(
            name=payload["name"], scope=payload["scope"], belongs_to=dis
        )
        ad = ad.dict()
        # Verify that both the datasets exist before attaching.
        ds = db_queries.fetch_dataset(session, payload["name"], payload["scope"])
        assert ds is not None, f"dataset {payload['name']}, {payload['scope']} not found"
        for bt in payload["belongs_to"]:
            belongs_to_ds = db_queries.fetch_dataset(session, bt["name"], bt["scope"])
            err_msg = f"dataset {bt['name']}, {bt['scope']} not found."
            assert belongs_to_ds is not None, err_msg
        # Now attach datasets
        status = db_queries.attach_datasets(session, **ad)

        ds = db_queries.fetch_dataset(session, payload["name"], payload["scope"])
        container: List[Any] = []
        db_queries.explore_children_datasets(ds, container)
        storage_map = db_queries.get_storage_map(session)
        for ds in container:
            if len(ds.files) > 0:
                rp, dp = policy_evaluator.determine_policies(session, ds, storage_map)
                policy_evaluator.apply_policy_to_dataset(
                    session, ds, rp, dp, storage_map
                )
        session.commit()
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error("Something errored out while attaching elements")
        log.error(str(e))
        return str(e)

create_dataset(session, payload, test_mode=False) async

Create a dataset.

Source code in datatrail/commit/commit_api.py

async def create_dataset(
    session: Union[AsyncSession, sessionmaker],
    payload: Union[common.DeletionPayload, common.ReplicationPayload],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Create a dataset."""
    try:
        conf = config.fetch_config()
        storage_elements = await get_storage_elements(session)
        if isinstance(storage_elements, str):
            return False
        ds = db_queries.fetch_dataset(session, payload["name"], payload["scope"])
        if ds:
            return True
        if "replication_policy" in payload:
            rp = datasets.ReplicationPolicy(**payload["replication_policy"])
            payload["replication_policy"] = rp
            replication_default = False
        else:
            payload["replication_policy"] = conf["default_replication_policy"]
            replication_default = True
        if "deletion_policy" in payload:
            dp: Any = [datasets.DeletionPolicy(**p) for p in payload["deletion_policy"]]
            # Add default policies for other sites depending on their type.
            sites = [d.storage_element for d in dp]
            for s in storage_elements:
                if s["name"] in sites:
                    continue
                policy = {
                    "storage_element": s["name"],
                    "delete_after_days": conf["default_deletion_policy"][
                        s["storage_type"]
                    ]["delete_after_days"],
                    "priority": conf["default_deletion_policy"][s["storage_type"]][
                        "priority"
                    ],
                }
                dp.append(datasets.DeletionPolicy(**policy))
            payload["deletion_policy"] = dp
            deletion_default = False
        else:
            # Add default policies for all sites depending on their type.
            ddp = []
            for s in storage_elements:
                dp = {
                    "storage_element": s["name"],
                    "delete_after_days": conf["default_deletion_policy"][
                        s["storage_type"]
                    ]["delete_after_days"],
                    "priority": conf["default_deletion_policy"][s["storage_type"]][
                        "priority"
                    ],
                }
                ddp.append(datasets.DeletionPolicy(**dp))
            payload["deletion_policy"] = ddp
            deletion_default = True
        ds = datasets.CreateDataset(**payload)
        payload = ds.dict()
        payload["replication_policy"]["default"] = replication_default
        for dp in payload["deletion_policy"]:
            dp["default"] = deletion_default
        vp = validate_policies(payload["replication_policy"], payload["deletion_policy"])
        if vp is not True:
            return vp
        status = db_queries.insert_dataset(session, **payload)
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

create_storage_element(session, payload) async

Create a storage element.

Source code in datatrail/commit/commit_api.py

async def create_storage_element(
    session: Union[AsyncSession, sessionmaker],
    payload: Union[common.DeletionPayload, common.ReplicationPayload],
) -> Union[bool, str]:
    """Create a storage element."""
    try:
        se = db_queries.fetch_storage_element(session, payload["name"])
        if se:
            return True
        cse = storage_element.CreateStorageElement(**payload)
        payload = cse.dict()
        status = db_queries.insert_storage_element(session, **payload)
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

detach_datasets(session, payload, test_mode=False) async

Detach two datasets.

Source code in datatrail/commit/commit_api.py

async def detach_datasets(
    session: Union[AsyncSession, sessionmaker],
    payload: Union[common.DeletionPayload, common.ReplicationPayload],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Detach two datasets."""
    try:
        dis = [datasets.DatasetIdentifier(**bt) for bt in payload["detach_from"]]
        ad = datasets.DetachDataset(
            name=payload["name"], scope=payload["scope"], detach_from=dis
        )
        ad = ad.dict()
        status = db_queries.detach_datasets(session, **ad)
        ds = db_queries.fetch_dataset(session, payload["name"], payload["scope"])
        container: List[Any] = []
        db_queries.explore_children_datasets(ds, container)
        storage_map = db_queries.get_storage_map(session)
        for ds in container:
            if len(ds.files) > 0:
                rp, dp = policy_evaluator.determine_policies(session, ds, storage_map)
                policy_evaluator.apply_policy_to_dataset(
                    session, ds, rp, dp, storage_map
                )
        session.commit()
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

register_files(session, payload, test_mode=False) async

Register files to a dataset.

Source code in datatrail/commit/commit_api.py

async def register_files(
    session: Union[AsyncSession, sessionmaker],
    payload: Union[common.DeletionPayload, common.ReplicationPayload],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Register files to a dataset."""
    try:
        storage_map = db_queries.get_storage_map(session)
        # Skip if files already exist for this dataset.
        conf = config.fetch_config()
        ds = db_queries.fetch_dataset(
            session, payload["dataset_name"], payload["dataset_scope"]
        )
        if ds.files:
            ds_file_names = [f.name for f in ds.files]
            fs_file_names = [f["name"] for f in payload["files"]]

            if set(fs_file_names) <= set(ds_file_names):
                return True
        for f in payload["files"]:
            # Check that name is path minus root directory.
            assert f["name"].startswith(
                "data/"
            ), "File names must start with data/, ie have site-specific root removed."
        # Continue otherwise.
        fs = [common.File(**f) for f in payload["files"]]

        pl = files.AttachFilesToDataset(
            dataset_name=payload["dataset_name"],
            dataset_scope=payload["dataset_scope"],
            files=fs,
        )
        pl = pl.dict()
        for f in pl["files"]:
            f["preferred_storage_elements"] = conf["default_replication_policy"][
                "preferred_storage_elements"
            ]
            f["storage_element_captured_at"] = storage_map[
                f["storage_element_captured_at"]
            ]
        status = db_queries.bulk_insert_files(session, **pl)
        session.commit()
        ds = db_queries.fetch_dataset(
            session, payload["dataset_name"], payload["dataset_scope"]
        )
        rp, dp = policy_evaluator.determine_policies(session, ds, storage_map)
        policy_evaluator.apply_policy_to_dataset(session, ds, rp, dp, storage_map)
        session.commit()
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

toggle_storage_element_state(session, name) async

Toggle state of storage element.

Source code in datatrail/commit/commit_api.py

async def toggle_storage_element_state(
    session: Union[AsyncSession, sessionmaker], name: str
) -> Union[bool, str]:
    """Toggle state of storage element."""
    try:
        status = db_queries.toggle_storage_element_state(session, name)
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

update_deletion_state(session, replica, deletion_state) async

Update the deletion state and time of the replica.

Source code in datatrail/commit/commit_api.py

async def update_deletion_state(
    session: Union[AsyncSession, sessionmaker],
    replica: models.FileReplica,
    deletion_state: str,
) -> models.FileReplica:
    """Update the deletion state and time of the replica."""
    replica.deletion_state = models.TransformationState[deletion_state]
    replica.deletion_state_updated_at = pytz.utc.localize(datetime.datetime.utcnow())
    return replica

update_deletion_status_full_dataset(session, payload, test_mode=False) async

Update the transformation status of the file replicas.

Source code in datatrail/commit/commit_api.py

async def update_deletion_status_full_dataset(
    session: Union[AsyncSession, sessionmaker],
    payload: Dict[str, str],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Update the transformation status of the file replicas."""
    try:
        dataset_name = payload["dataset_name"]
        dataset_scope = payload["dataset_scope"]
        storage_name = payload["storage_name"]
        status = payload["status"]
        ds = db_queries.fetch_dataset(session, dataset_name, dataset_scope)
        storage_map = db_queries.get_storage_map(session)
        files = set()
        for f in ds.files:
            for fr in f.file_replicas:
                if fr.storage_id == storage_map[storage_name]:
                    fr = await update_deletion_state(session, fr, status)
                    if status == "completed" and fr.belongs_to.name not in files:
                        f.num_available_replicas = len(
                            [
                                r
                                for r in f.file_replicas
                                if r.deletion_state
                                != models.TransformationState.completed
                            ]
                        )
                        files.add(f.name)
                    break
        session.commit()
        return True
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

update_replication_policy(session, scope, name, replication_policy, test_mode=False) async

Update the replication policy of a dataset.

Source code in datatrail/commit/commit_api.py

async def update_replication_policy(
    session: Union[AsyncSession, sessionmaker],
    scope: str,
    name: str,
    replication_policy: datasets.ReplicationPolicy,
    test_mode: bool = False,
) -> Union[bool, str]:
    """Update the replication policy of a dataset."""
    try:
        rp = datasets.ReplicationPolicy(**replication_policy)
        urp = datasets.UpdateReplicationPolicy(name=name, scope=scope, policy=rp)
        urp = urp.dict()
        dataset = await get_dataset(session, scope, name, test_mode=test_mode)
        if isinstance(dataset, str):
            return "Could not find dataset"
        vp = validate_policies(urp["policy"], dataset["deletion_policy"])
        if vp is not True:
            return vp
        urp["policy"]["default"] = False
        status = db_queries.update_replication_policy(session, **urp)
        this_dataset = db_queries.fetch_dataset(session, name, scope)
        container: List[Any] = []
        db_queries.explore_children_datasets(this_dataset, container)
        storage_map = db_queries.get_storage_map(session)
        for ds in container:
            if len(ds.files) > 0:
                rp_ds, dp_ds = policy_evaluator.determine_policies(
                    session, ds, storage_map
                )
                policy_evaluator.apply_policy_to_dataset(
                    session, ds, rp_ds, dp_ds, storage_map
                )
        session.commit()
        return status
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

update_replication_state(session, replica, replication_state) async

Update the replication state and time of the replica.

Source code in datatrail/commit/commit_api.py

async def update_replication_state(
    session: Union[AsyncSession, sessionmaker],
    replica: models.FileReplica,
    replication_state: str,
) -> models.FileReplica:
    """Update the replication state and time of the replica."""
    replica.replication_state = models.TransformationState[replication_state]
    replica.replication_state_updated_at = pytz.utc.localize(datetime.datetime.utcnow())
    if replication_state == "completed":
        pse = deepcopy(replica.belongs_to.preferred_storage_elements)
        if replica.replicate_to is not None:
            storage_map = db_queries.get_storage_map(session)
            if storage_map[replica.replicate_to] in pse:
                pse.remove(storage_map[replica.replicate_to])
            if replica.storage_id in pse:
                pse.remove(replica.storage_id)
            if pse == []:
                replica.replicate_to = None
            else:
                replica.replicate_to = storage_map[pse[0]]
    return replica

update_transformation_status(session, payload, test_mode=False) async

Update the transformation status of the file replica.

Source code in datatrail/commit/commit_api.py

async def update_transformation_status(
    session: Union[AsyncSession, sessionmaker],
    payload: Dict[int, common.TransformationStatus],
    test_mode: bool = False,
) -> Union[bool, str]:
    """Update the transformation status of the file replica."""
    try:
        files = set()
        for replica_id, transformation_status in payload.items():
            transformation_status = common.TransformationStatus(
                **transformation_status
            ).dict()
            replica = db_queries.fetch_file_replica_by_id(session, replica_id)
            if transformation_status["transformation"] == "replication":
                replica = await update_replication_state(
                    session, replica, transformation_status["status"]
                )
            else:
                # transformation is deletion
                replica = await update_deletion_state(
                    session, replica, transformation_status["status"]
                )
                if (
                    transformation_status["status"] == "completed"
                    and replica.belongs_to.name not in files
                ):
                    f = replica.belongs_to
                    f.num_available_replicas = len(
                        [
                            fr
                            for fr in f.file_replicas
                            if fr.deletion_state != models.TransformationState.completed
                        ]
                    )
                    files.add(f.name)
        session.commit()
        return True
    except Exception as e:  # pragma: no cover
        session.rollback()
        log.error(str(e))
        return str(e)

validate_policies(replication_policy, deletion_policy)

Validate policies.

Ensure that each preferred storage element has a deletion policy.

Source code in datatrail/commit/commit_api.py

def validate_policies(
    replication_policy: datasets.ReplicationPolicy,
    deletion_policy: List[Any],
) -> Union[bool, str]:
    """Validate policies.

    Ensure that each preferred storage element has a deletion policy.
    """
    valid: Union[bool, str] = True
    pse = replication_policy["preferred_storage_elements"]
    deletion_storage_elements = [d["storage_element"] for d in deletion_policy]
    for se in pse:
        if se not in deletion_storage_elements:
            valid = f"Error! No deletion policy found for {se}."
    return valid