def test_headers():
header = Headers({"Header-Key": "HeaderValue"})
assert header == {"header-key": "HeaderValue"}
header["Header-Key-2"] = "HeaderValue2"
assert header == {
"header-key": "HeaderValue",
"header-key-2": "HeaderValue2"
}
header.update({"Header-Key-3": "HeaderValue3"})
assert header == {
"header-key": "HeaderValue",
"header-key-2": "HeaderValue2",
"header-key-3": "HeaderValue3",
}
header.update({"Header-Key": "ChangedHeaderValue"})
assert header == {
"header-key": "ChangedHeaderValue",
"header-key-2": "HeaderValue2",
"header-key-3": "HeaderValue3",
}
assert header.get("Header-Key") == "ChangedHeaderValue"
assert header.get("HeAdEr-KeY") == "ChangedHeaderValue"
assert header.get("header-key") == "ChangedHeaderValue"
assert header.get("HEADER-KEY") == "ChangedHeaderValue"
assert header["Header-Key"] == "ChangedHeaderValue"
assert header["HeAdEr-KeY"] == "ChangedHeaderValue"
assert header["header-key"] == "ChangedHeaderValue"
assert header["HEADER-KEY"] == "ChangedHeaderValue"
def set_datasets_provenance(self, provenance):
# type(Dict[Any, Any]) -> Dict[Any, Any]
""" Sets the Dataset Provenance
Args:
provenance (dict): Provenance dictionary
Returns:
dict: Provenance response dictionary.
Example:
.. code-block:: python
from crux import Crux
conn = Crux()
provenance = {
"dataset_id":[
{
"workflow_id": "test_id",
"pipeline_ids": ["test_id_1","test_id_2"],
"cron_spec": "0 0 1 1 0"
}
]
}
response = conn.set_datasets_provenance(provenance=provenance)
"""
headers = Headers({"accept": "application/json"})
response = self.api_client.api_call(
"POST", ["datasets", "provenance"], headers=headers, json=provenance
)
return response.json()
def get_resource(self, id): # id is by design pylint: disable=redefined-builtin
# type: (str) -> Union[File, Folder]
"""Fetches the Resource by ID.
Any supported resource can be fetched. The object returned will be an instance
of the correct subclass, for example a ``crux.models.File`` instance will be
returned for file resources.
Args:
id (str): Resource ID which is to be fetched.
Returns:
crux.models.Resource: Resource or its Child Object.
"""
headers = Headers(
{"accept": "application/json"}
) # type: MutableMapping[Text, Text]
response = self.api_client.api_call("GET", ["resources", id], headers=headers)
raw_resource = response.json()
resource = get_resource_object(
resource_type=raw_resource.get("type"),
data=raw_resource,
connection=self.api_client,
)
return resource
def add_labels(self, labels_dict):
# type: (dict) -> bool
"""Adds multiple labels to Resource.
Args:
label_dict (dict): Labels (key/value pairs) to add to the Resource.
Returns:
bool: True if the labels were added, False otherwise.
"""
headers = Headers({"content-type": "application/json", "accept": "application/json"})
labels_list = []
for label_key, label_value in labels_dict.items():
if label_key is not None and label_value is not None:
label_key = label_key.value if isinstance(label_key, Enum) else label_key
labels_list.append(
{"labelKey": str(label_key), "labelValue": str(label_value)}
)
data = {"labels": labels_list}
response_result = self.connection.api_call(
"PUT",
["datasets", self.dataset_id, "resources", self.id, "labels"],
headers=headers,
json=data,
)
if response_result:
# Sync the latest data from API to prevent inconsistency
self.refresh()
return True
def add_label(self, label_key, label_value):
# type: (str, str) -> bool
"""Adds label to Resource.
Args:
label_key (str): Label Key for Resource.
label_value (str): Label Value for Resource.
Returns:
bool: True if label is added, False otherwise.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
response_result = self.connection.api_call(
"PUT",
[
"datasets",
self.dataset_id,
"resources",
self.id,
"labels",
label_key,
label_value,
],
headers=headers,
)
if response_result:
# Sync the latest data from API to prevent inconsistency
self.refresh()
return True
def delete_label(self, label_key):
# type: (str) -> bool
"""Deletes label from Resource.
Args:
label_key (str): Label Key for Resource.
Returns:
bool: True if label is deleted, False otherwise.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
response_result = self.connection.api_call(
"DELETE",
[
"datasets", self.dataset_id, "resources", self.id, "labels",
label_key
],
headers=headers,
)
if response_result:
# Sync the latest data from API to prevent inconsistency
self.refresh()
return True
def list_datasets(self, owned=True, subscribed=True):
# type: (bool, bool) -> List[Dataset]
"""Fetches a list of owned and/or subscribed Datasets.
Args:
owned (bool): Show datasets owned by the caller. Defaults to True.
subscribed (bool): Show datasets the user has a subscription. Defaults to True.
Returns:
list(:obj:`crux.models.Dataset`): List of Dataset objects.
"""
dataset_list = []
# Prefer domainV2 for data source
headers = Headers({"accept": "application/json"
}) # type: MutableMapping[Text, Text]
pagesize = 100
params = {"limit": pagesize}
retrieved = 0
while True:
params["offset"] = retrieved
try:
resp = self.api_client.api_call(
"GET",
["v2", "client", "subscriptions", "view", "summary"],
params=params,
model=None,
headers=headers,
).json()
except CruxAPIError as err:
log.debug("Get subscriptions failed: %s", err)
break
for dataset in resp:
dataset["name"] = dataset["datasetName"]
obj = Dataset.from_dict(dataset, connection=self.api_client)
dataset_list.append(obj)
respcnt = len(resp)
retrieved += respcnt
if respcnt < pagesize:
break
if dataset_list:
return dataset_list
# Try legacy tables
datasets = self._call_drives_my()
if owned:
for dataset in datasets["owned"]:
obj = Dataset.from_dict(dataset, connection=self.api_client)
dataset_list.append(obj)
if subscribed:
for dataset in datasets["subscriptions"]:
obj = Dataset.from_dict(dataset, connection=self.api_client)
dataset_list.append(obj)
return dataset_list
def _call_drives_my(self):
headers = Headers({"accept": "application/json"
}) # type: MutableMapping[Text, Text]
response = self.api_client.api_call("GET", ["drives", "my"],
model=None,
headers=headers)
return response.json()
def delete(self):
# type: () -> bool
"""Deletes Resource from Dataset.
Returns:
bool: True if it is deleted.
"""
headers = Headers({"content-type": "application/json", "accept": "application/json"})
return self.connection.api_call("DELETE", ["resources", self.id], headers=headers)
def list_permissions(self):
# type: () -> List[Permission]
"""Lists the permission on the resource.
Returns:
list (:obj:`crux.models.Permission`): List of Permission Objects.
"""
headers = Headers({"accept": "application/json"})
return self.connection.api_call(
"GET", ["resources", self.id, "permissions"], model=Permission, headers=headers,
)
def list_public_datasets(self):
# type: () -> List[Dataset]
"""Fetches a list of public Datasets.
Returns:
list (:obj:`crux.models.Dataset`): List of Dataset objects.
"""
headers = Headers({"accept": "application/json"
}) # type: MutableMapping[Text, Text]
return self.api_client.api_call("GET", ["datasets", "public"],
model=Dataset,
headers=headers)
def whoami(self):
# type: () -> Identity
"""Returns the Identity of Current User.
Returns:
crux.models.Identity: Identity object.
"""
headers = Headers({"accept": "application/json"
}) # type: Optional[MutableMapping[Text, Text]]
return self.api_client.api_call("GET", ["identities", "whoami"],
model=Identity,
headers=headers)
def _get_folder(self):
# type: () -> str
"""Fetches the folder of the resource.
Returns:
str: Folder name of the resource.
"""
headers = Headers({"content-type": "application/json", "accept": "application/json"})
response = self.connection.api_call(
"GET", ["resources", self.id, "folderpath"], headers=headers
)
return response.json().get("path")
def add_permission( # It is by design pylint: disable=arguments-differ
self, identity_id, permission, recursive=False
):
# type: (str, str, bool) -> Union[bool, Permission]
"""Adds permission to the Folder resource.
Args:
identity_id (str): Identity Id to be set.
permission (str): Permission to be set.
recursive (bool): If recursive is set to True, it will recursive apply
permission to all resources under the folder resource.
Returns:
bool or crux.models.Permission: If recursive is set then it returns True.
If recursive is unset then it returns Permission object.
"""
headers = Headers(
{"content-type": "application/json", "accept": "application/json"}
)
body = {
"identityId": identity_id,
"permission": permission,
"action": "add",
"resourceIds": [self.id],
}
if recursive:
log.debug(
"Adding permission %s for %s in recursive mode to resource %s",
permission,
identity_id,
self.id,
)
return self.connection.api_call(
"POST", ["permissions", "bulk"], headers=headers, json=body
)
else:
log.debug(
"Adding permission %s for %s to resource %s",
permission,
identity_id,
self.id,
)
return self.connection.api_call(
"PUT",
["permissions", self.id, identity_id, permission],
model=Permission,
headers=headers,
)
def delete_permission(self, identity_id, permission):
# type: (str, str) -> bool
"""Deletes permission from the resource.
Args:
identity_id (str): Identity Id for the deletion.
permission (str): Permission for the deletion.
Returns:
bool: True if it is able to delete it.
"""
headers = Headers({"content-type": "application/json", "accept": "application/json"})
return self.connection.api_call(
"DELETE", ["permissions", self.id, identity_id, permission], headers=headers
)
def _download(self, file_obj, media_type, chunk_size=DEFAULT_CHUNK_SIZE):
if media_type is not None:
headers = Headers({"accept": media_type})
else:
headers = None
data = self.connection.api_call(
"GET", ["resources", self.id, "content"], headers=headers, stream=True
)
for chunk in data.iter_content(chunk_size=chunk_size):
file_obj.write(chunk)
data.close()
return True
def get_job(self, job_id):
# type: (str) -> Job
"""Fetches the Job.
Args:
job_id (str): Job ID which is to be fetched.
Returns:
crux.models.Job: Job object.
"""
headers = Headers({"accept": "application/json"
}) # type: MutableMapping[Text, Text]
return self.api_client.api_call("GET", ["jobs", job_id],
model=Job,
headers=headers)
def get_dataset(self, id): # id name is by design pylint: disable=redefined-builtin
# type: (str) -> Dataset
"""Fetches the Dataset.
Args:
id (str): Dataset ID which is to be fetched.
Returns:
crux.models.Dataset: Dataset object
"""
headers = Headers({"accept": "application/json"
}) # type: MutableMapping[Text, Text]
return self.api_client.api_call("GET", ["datasets", id],
model=Dataset,
headers=headers)
def delete_permission( # It is by design pylint: disable=arguments-differ
self, identity_id, permission, recursive=False
):
# type: (str, str, bool) -> bool
"""Deletes permission from Folder resource.
Args:
identity_id (str): Identity Id for the deletion.
permission (str): Permission for deletion.
recursive (bool): If recursive is set to True, it will recursively delete
permission from all resources under the folder resource.
Defaults to False.
Returns:
bool: True if it is able to delete it.
"""
headers = Headers(
{"content-type": "application/json", "accept": "application/json"}
)
body = {
"identityId": identity_id,
"permission": permission,
"action": "delete",
"resourceIds": [self.id],
}
if recursive:
log.debug(
"Deleting permission %s for %s in recursive mode from resource %s",
permission,
identity_id,
self.id,
)
return self.connection.api_call(
"POST", ["permissions", "bulk"], headers=headers, json=body
)
else:
log.debug(
"Deleting permission %s for %s in from resource %s",
permission,
identity_id,
self.id,
)
return self.connection.api_call(
"DELETE",
["permissions", self.id, identity_id, permission],
headers=headers,
)
def refresh(self):
"""Refresh Resource model from API backend.
Returns:
bool: True, if it is able to refresh the model,
False otherwise.
"""
# type () -> bool
headers = Headers({"content-type": "application/json", "accept": "application/json"})
resource_object = self.connection.api_call(
"GET", ["resources", self.id], headers=headers, model=Resource
)
self.raw_model = resource_object.raw_model
return True
def update(self, name=None, description=None, tags=None, provenance=None):
# type: (str, str, List[str], str) -> bool
"""Updates the metadata for Resource.
Args:
name (str): Name of resource. Defaults to None.
description (str): Description of the resource. Defaults to None.
tags (:obj:`list` of :obj:`str`): List of tags. Defaults to None.
provenance (str): Provenance for a resource. Defaults to None.
Returns:
bool: True, if resource is updated.
Raises:
ValueError: It is raised if name, description or tags are unset.
TypeError: It is raised if tags are not of type List.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
body = {} # type: Dict[str, Union[str, List, Dict]]
if name is not None:
self.raw_model["name"] = name
if description is not None:
self.raw_model["description"] = description
if tags is not None:
self.raw_model["tags"] = tags
if provenance is not None:
self.raw_model["provenance"] = provenance
body = self.raw_model
log.debug("Body %s", body)
resource_object = self.connection.api_call("PUT",
["resources", self.id],
headers=headers,
json=body,
model=Resource)
self.raw_model = resource_object.raw_model
log.debug("Updated dataset %s with content %s", self.id,
self.raw_model)
return True
def get_stitch_job(self, job_id):
# type: (str) -> StitchJob
"""Stitch Job Details.
Args:
job_id (str): Job ID of the Stitch Job.
Returns:
crux.models.StitchJob: StitchJob object.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
return self.connection.api_call("GET", ["datasets", "stitch", job_id],
headers=headers,
model=StitchJob)
def delete_label(self, label_key):
# type: (str) -> bool
"""Deletes label from Dataset.
Args:
label_key (str): Label Key for Dataset.
Returns:
bool: True if labels are deleted.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
return self.connection.api_call(
"DELETE", ["datasets", self.id, "labels", label_key],
headers=headers)
def add_permission(self, identity_id, permission):
# type: (str, str) -> Union[bool, Permission]
"""Adds permission to the resource.
Args:
identity_id: Identity Id to be set.
permission: Permission to be set.
Returns:
crux.models.Permission: Permission Object.
"""
headers = Headers({"content-type": "application/json", "accept": "application/json"})
return self.connection.api_call(
"PUT",
["permissions", self.id, identity_id, permission],
model=Permission,
headers=headers,
)
def create_folder(self, path, folder="/", tags=None, description=None):
# type: (str, str, List[str], str) -> Folder
"""Creates Folder resource in Dataset.
Args:
path (str): Path of the Folder resource.
folder (str): Parent folder of the Folder resource.
Defaults to /.
tags (:obj:`list` of :obj:`str`): Tags of the Folder resource.
Defaults to None.
description (str): Description of the Folder resource.
Defaults to None.
Returns:
crux.models.Folder: Folder Object.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
tags = tags if tags else []
file_name, folder = split_posixpath_filename_dirpath(path)
raw_model = {
"name": file_name,
"type": "folder",
"tags": tags,
"description": description,
"folder": folder,
}
folder_resource = Folder(raw_model=raw_model)
return self.connection.api_call(
"POST",
["datasets", self.id, "resources"],
json=folder_resource.raw_model,
model=Folder,
headers=headers,
)
def create(self):
# type: () -> bool
"""Creates the Dataset.
Returns:
bool: True if dataset is created.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
dataset_object = self.connection.api_call("POST", ["datasets"],
json=self.raw_model,
model=Dataset,
headers=headers)
self.raw_model = dataset_object.raw_model
return True
def get_label(self, label_key):
# type: (str) -> Label
"""Gets label value of Dataset.
Args:
label_key (str): Label Key for Dataset.
Returns:
crux.models.Label: Label Object.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
return self.connection.api_call(
"GET",
["datasets", self.id, "labels", label_key],
headers=headers,
model=Label,
)
def add_label(self, label_key, label_value):
# type: (str, str) -> bool
"""Adds label to Dataset.
Args:
label_key (str): Label Key for Dataset.
label_value (str): Label Value for Dataset.
Returns:
bool: True if labels are added.
"""
headers = Headers({
"content-type": "application/json",
"accept": "application/json"
})
return self.connection.api_call(
"PUT",
["datasets", self.id, "labels", label_key, label_value],
headers=headers,
)
def get_delivery(self, delivery_id):
# type: (str) -> Delivery
"""Gets Delivery object.
Args:
delivery_id (str): Delivery ID.
Returns:
crux.models.Delivery: Delivery Object.
Raises:
ValueError: If delivery_id value is invalid.
"""
headers = Headers({"accept": "application/json"})
if not DELIVERY_ID_REGEX.match(delivery_id):
raise ValueError("Value of delivery_id is invalid")
return self.connection.api_call("GET",
["deliveries", self.id, delivery_id],
headers=headers,
model=Delivery)
def get_ingestions(self, start_date=None, end_date=None):
# type: (str, str) -> Iterator[Ingestion]
"""Gets Ingestions.
Args:
start_date (str): ISO format start time.
end_date (str): ISO format end time.
Returns:
crux.models.Delivery: Delivery Object.
"""
headers = Headers({"accept": "application/json"})
params = {}
params["start_date"] = start_date
params["end_date"] = end_date
response = self.connection.api_call("GET",
["deliveries", self.id, "ids"],
headers=headers,
params=params)
all_deliveries = response.json()
ingestion_map = defaultdict(set) # type: DefaultDict[str, Set]
for delivery in all_deliveries:
if not DELIVERY_ID_REGEX.match(delivery):
raise ValueError("Value of delivery_id is invalid")
ingestion_id, version_id = delivery.split(".")
ingestion_map[ingestion_id].add(int(version_id))
for ingestion_id in ingestion_map:
obj = Ingestion.from_dict({
"ingestionId": ingestion_id,
"versions": ingestion_map[ingestion_id],
"datasetId": self.id,
})
obj.connection = self.connection
yield obj