def _job_to_proto(spark_job: SparkJob) -> JobProto:
job = JobProto()
job.id = spark_job.get_id()
job.log_uri = cast(str, spark_job.get_log_uri() or "")
status = spark_job.get_status()
if status == SparkJobStatus.COMPLETED:
job.status = JobStatus.JOB_STATUS_DONE
elif status == SparkJobStatus.IN_PROGRESS:
job.status = JobStatus.JOB_STATUS_RUNNING
elif status == SparkJobStatus.FAILED:
job.status = JobStatus.JOB_STATUS_ERROR
elif status == SparkJobStatus.STARTING:
job.status = JobStatus.JOB_STATUS_PENDING
else:
raise ValueError(f"Invalid job status {status}")
if isinstance(spark_job, RetrievalJob):
job.type = JobType.RETRIEVAL_JOB
job.retrieval.output_location = spark_job.get_output_file_uri(
block=False)
elif isinstance(spark_job, BatchIngestionJob):
job.type = JobType.BATCH_INGESTION_JOB
job.batch_ingestion.table_name = spark_job.get_feature_table()
elif isinstance(spark_job, StreamIngestionJob):
job.type = JobType.STREAM_INGESTION_JOB
job.stream_ingestion.table_name = spark_job.get_feature_table()
else:
raise ValueError(f"Invalid job type {job}")
job.start_time.FromDatetime(spark_job.get_start_time())
return job
def GetJob(self, request, context):
self._call_count["GetJob"] += 1
return GetJobResponse(job=JobProto(
id="test",
type=JobType.RETRIEVAL_JOB,
status=self._job_statuses.pop(),
retrieval=JobProto.RetrievalJobMeta(output_location="foo"),
))