def _get_executions_by_pipeline_name(
self, pipeline_name: Text) -> List[metadata_store_pb2.Execution]:
"""Helper function returns executions under a given pipeline name."""
# step 1: get context id by context name
request = metadata_store_service_pb2.GetContextByTypeAndNameRequest(
type_name='pipeline', context_name=pipeline_name)
context_id = self._stub.GetContextByTypeAndName(request).context.id
# step 2: get executions by context id
request = metadata_store_service_pb2.GetExecutionsByContextRequest(
context_id=context_id)
return self._stub.GetExecutionsByContext(request).executions
def _get_artifacts_with_type_and_pipeline(
self, type_name: Text,
pipeline_name: Text) -> List[metadata_store_pb2.Artifact]:
"""Helper function returns artifacts of specified pipeline and type."""
# 1. Find the pipeline context according to its name.
request = metadata_store_service_pb2.GetContextByTypeAndNameRequest(
type_name=_CONTEXT_TYPE_PIPELINE, context_name=pipeline_name)
pipeline_context = self._stub.GetContextByTypeAndName(request)
# 2. Find the artifacts associated with the pipeline context.
request = metadata_store_service_pb2.GetArtifactsByContextRequest(
context_id=pipeline_context.context.id)
artifacts_response = self._stub.GetArtifactsByContext(request)
# 3. Find the specified artifact type id.
artifact_type_request = metadata_store_service_pb2.GetArtifactTypeRequest(
type_name=type_name)
artifact_type = self._stub.GetArtifactType(
artifact_type_request).artifact_type
# 4. Filter the returned artifacts according to their types and return.
return [
artifact for artifact in artifacts_response.artifacts
if artifact.type_id == artifact_type.id
]
def get_context_by_type_and_name(
self, type_name: Text,
context_name: Text) -> Optional[metadata_store_pb2.Context]:
"""Get the context of the given type and context name.
The API fails if more than one contexts are found.
Args:
type_name: The context type name to look for.
context_name: The context name to look for.
Returns:
The Context matching the type and context name.
None if no matched Context found.
"""
request = metadata_store_service_pb2.GetContextByTypeAndNameRequest()
request.type_name = type_name
request.context_name = context_name
response = metadata_store_service_pb2.GetContextByTypeAndNameResponse()
self._call('GetContextByTypeAndName', request, response)
if not response.HasField('context'):
return None
return response.context
