Bugfix: Python SDK listing of ingestion job fails for featureset reference filter (#861)

terryyylim · web-flow · commit bc7f3fb04b5d · 2020-07-07T16:27:03.000+08:00
* Fix listing of ingestion jobs for featureset ref

* Fix lint

* Shift filter to java method
diff --git a/core/src/main/java/feast/core/dao/JobRepository.java b/core/src/main/java/feast/core/dao/JobRepository.java
@@ -40,6 +40,9 @@ public interface JobRepository extends JpaRepository<Job, String> {
 
   List<Job> findByFeatureSetJobStatusesIn(List<FeatureSetJobStatus> featureSetsJobStatuses);
 
+  List<Job> findByFeatureSetJobStatusesFeatureSetNameAndFeatureSetJobStatusesFeatureSetProjectName(
+      String featureSetName, String featureSetProject);
+
   // find jobs that have at least one store with given name
   List<Job> findByJobStoresIdStoreName(String storeName);
 }
diff --git a/core/src/main/java/feast/core/service/JobService.java b/core/src/main/java/feast/core/service/JobService.java
@@ -23,11 +23,9 @@
 import feast.core.log.Action;
 import feast.core.log.AuditLogger;
 import feast.core.log.Resource;
-import feast.core.model.FeatureSet;
 import feast.core.model.Job;
 import feast.core.model.JobStatus;
 import feast.proto.core.CoreServiceProto.ListFeatureSetsRequest;
-import feast.proto.core.CoreServiceProto.ListFeatureSetsResponse;
 import feast.proto.core.CoreServiceProto.ListIngestionJobsRequest;
 import feast.proto.core.CoreServiceProto.ListIngestionJobsResponse;
 import feast.proto.core.CoreServiceProto.RestartIngestionJobRequest;
@@ -115,20 +113,17 @@ public ListIngestionJobsResponse listJobs(ListIngestionJobsRequest request)
         if (filter.hasFeatureSetReference()) {
           // find a matching featuresets for reference
           FeatureSetReference fsReference = filter.getFeatureSetReference();
-          ListFeatureSetsResponse response =
-              this.specService.listFeatureSets(this.toListFeatureSetFilter(fsReference));
-          List<FeatureSet> featureSets =
-              response.getFeatureSetsList().stream()
-                  .map(FeatureSet::fromProto)
-                  .collect(Collectors.toList());
 
           // find jobs for the matching featuresets
           Collection<Job> matchingJobs =
-              this.jobRepository.findByFeatureSetJobStatusesIn(
-                  featureSets.stream()
-                      .flatMap(fs -> fs.getJobStatuses().stream())
-                      .collect(Collectors.toList()));
-          List<String> jobIds = matchingJobs.stream().map(Job::getId).collect(Collectors.toList());
+              this.jobRepository
+                  .findByFeatureSetJobStatusesFeatureSetNameAndFeatureSetJobStatusesFeatureSetProjectName(
+                      fsReference.getName(), fsReference.getProject());
+          List<String> jobIds =
+              matchingJobs.stream()
+                  .filter(job -> job.getStatus().equals(JobStatus.RUNNING))
+                  .map(Job::getId)
+                  .collect(Collectors.toList());
           matchingJobIds = this.mergeResults(matchingJobIds, jobIds);
         }
       }
diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py
@@ -749,16 +749,21 @@ def list_ingest_jobs(
             List of IngestJobs matching the given filters
         """
         # construct list request
-        feature_set_ref = None
+        feature_set_ref_proto = None
+        if feature_set_ref:
+            feature_set_ref_proto = feature_set_ref.to_proto()
         list_filter = ListIngestionJobsRequest.Filter(
-            id=job_id, feature_set_reference=feature_set_ref, store_name=store_name,
+            id=job_id,
+            feature_set_reference=feature_set_ref_proto,
+            store_name=store_name,
         )
         request = ListIngestionJobsRequest(filter=list_filter)
         # make list request & unpack response
-        response = self._core_service_stub.ListIngestionJobs(request)  # type: ignore
+        response = self._core_service.ListIngestionJobs(request)  # type: ignore
         ingest_jobs = [
-            IngestJob(proto, self._core_service_stub) for proto in response.jobs  # type: ignore
+            IngestJob(proto, self._core_service) for proto in response.jobs  # type: ignore
         ]
+
         return ingest_jobs
 
     def restart_ingest_job(self, job: IngestJob):

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,9 @@ public interface JobRepository extends JpaRepository<Job, String> {`
`40`	`40`
`41`	`41`	`List<Job> findByFeatureSetJobStatusesIn(List<FeatureSetJobStatus> featureSetsJobStatuses);`
`42`	`42`
	`43`	`+ List<Job> findByFeatureSetJobStatusesFeatureSetNameAndFeatureSetJobStatusesFeatureSetProjectName(`
	`44`	`+ String featureSetName, String featureSetProject);`
	`45`	`+`
`43`	`46`	`// find jobs that have at least one store with given name`
`44`	`47`	`List<Job> findByJobStoresIdStoreName(String storeName);`
`45`	`48`	`}`