CassandraJobStore implements JobStore { @Override public Observable<Pair<List<Job<?>>, Integer>> retrieveJobs() { Observable result = Observable.fromCallable(() -> { List<String> jobIds = activeJobIdsBucketManager.getItems(); return jobIds.stream().map(retrieveActiveJobStatement::bind).map(this::execute).collect(Collectors.toList()); }).flatMap(observables -> Observable.merge(observables, getConcurrencyLimit()).flatMapIterable(resultSet -> { List<Row> allRows = resultSet.all(); if (allRows.isEmpty()) { logger.debug("Job id with no record"); return Collections.emptyList(); } return allRows.stream() .map(row -> row.getString(0)) .map(value -> { String effectiveValue; if (fitBadDataInjection.isPresent()) { effectiveValue = fitBadDataInjection.get().afterImmediate(JobStoreFitAction.ErrorKind.CorruptedRawJobRecords.name(), value); } else { effectiveValue = value; } Job<?> job; try { job = ObjectMappers.readValue(mapper, effectiveValue, Job.class); } catch (Exception e) { logger.error("Cannot map serialized job data to Job class: {}", effectiveValue, e); return Either.ofError(e); } if (job.getJobDescriptor().getDisruptionBudget() == null) { titusRuntime.getCodeInvariants().inconsistent("jobWithNoDisruptionBudget: jobId=%s", job.getId()); job = JobFunctions.changeDisruptionBudget(job, DisruptionBudget.none()); } if (job.getJobDescriptor().getExtensions() instanceof ServiceJobExt) { Job<ServiceJobExt> serviceJob = (Job<ServiceJobExt>) job; ServiceJobExt ext = serviceJob.getJobDescriptor().getExtensions(); if (ext.getMigrationPolicy() == null) { titusRuntime.getCodePointTracker().markReachable("Corrupted task migration record in Cassandra: " + job.getId()); ServiceJobExt fixedExt = ext.toBuilder().withMigrationPolicy(SystemDefaultMigrationPolicy.newBuilder().build()).build(); logger.warn("Service job with no migration policy defined. Setting system default: {}", job.getId()); job = serviceJob.toBuilder().withJobDescriptor( serviceJob.getJobDescriptor().toBuilder().withExtensions(fixedExt).build() ).build(); } } if (!fitBadDataInjection.isPresent()) { return Either.ofValue(job); } Job<?> effectiveJob = fitBadDataInjection.get().afterImmediate(JobStoreFitAction.ErrorKind.CorruptedJobRecords.name(), job); return Either.ofValue(effectiveJob); }) .collect(Collectors.toList()); })).toList().map(everything -> { List<Job> goodJobs = (List<Job>) everything.stream().filter(Either::hasValue).map(Either::getValue).collect(Collectors.toList()); int errors = everything.size() - goodJobs.size(); return Pair.of(goodJobs, errors); }); return result; } @Inject  CassandraJobStore(CassandraStoreConfiguration configuration,
                             Session session,
                             TitusRuntime titusRuntime);   CassandraJobStore(CassandraStoreConfiguration configuration,
                      Session session,
                      TitusRuntime titusRuntime,
                      ObjectMapper mapper,
                      int initialBucketCount,
                      int maxBucketSize); @Override Completable init(); @Override Observable<Pair<List<Job<?>>, Integer>> retrieveJobs(); @Override Observable<Job<?>> retrieveJob(String jobId); @Override Completable storeJob(Job job); @Override Completable updateJob(Job job); @Override Completable deleteJob(Job job); @Override Observable<Pair<List<Task>, Integer>> retrieveTasksForJob(String jobId); @Override Observable<Task> retrieveTask(String taskId); @Override Completable storeTask(Task task); @Override Completable updateTask(Task task); @Override Completable replaceTask(Task oldTask, Task newTask); @Override Completable moveTask(Job jobFrom, Job jobTo, Task taskAfter); @Override Completable deleteTask(Task task); @Override Observable<Job<?>> retrieveArchivedJob(String jobId); @Override Observable<Task> retrieveArchivedTasksForJob(String jobId); @Override Observable<Task> retrieveArchivedTask(String taskId);  }
@Test public void testRetrieveJobs() { Session session = cassandraCqlUnit.getSession(); JobStore bootstrappingStore = getJobStore(session); Job<BatchJobExt> job = createBatchJobObject(); bootstrappingStore.storeJob(job).await(); JobStore store = getJobStore(session); store.init().await(); Pair<List<Job<?>>, Integer> jobsAndErrors = store.retrieveJobs().toBlocking().first(); assertThat(jobsAndErrors.getLeft()).hasSize(1); assertThat(jobsAndErrors.getRight()).isEqualTo(0); assertThat(jobsAndErrors.getLeft().get(0)).isEqualTo(job); try { store.retrieveArchivedJob(job.getId()).toBlocking().first(); fail("Should not return active job"); } catch (JobStoreException e) { assertThat(e.getErrorCode()).isEqualTo(JobStoreException.ErrorCode.JOB_DOES_NOT_EXIST); } }