Coverage for rest_api/data_entry/sample_entry_job.py: 13%

1from datetime import datetime 

2import pathlib 

3import pickle 

4import shutil 

5import traceback 

6import zipfile 

7 

8from celery import group 

9from celery import shared_task 

10from django.core.cache import cache 

11from django.core.exceptions import FieldDoesNotExist 

12from django.db import DataError 

13from django.db import transaction 

14from django.db.models import Q 

15from django.utils import timezone 

16from line_profiler import LineProfiler 

17import pandas as pd 

18 

19from rest_api import models 

20from rest_api.data_entry.annotation_import import AnnotationImport 

21from rest_api.data_entry.sample_import import SonarImport 

22from rest_api.models import Alignment 

23from rest_api.models import AminoAcidMutation 

24from rest_api.models import AnnotationType 

25from rest_api.models import FileProcessing 

26from rest_api.models import ImportLog 

27from rest_api.models import NucleotideMutation 

28from rest_api.models import ProcessingJob 

29from rest_api.models import Sample 

30from rest_api.models import Sequence 

31from rest_api.serializers import Sample2PropertyBulkCreateOrUpdateSerializer 

32from rest_api.utils import parse_date 

33from rest_api.utils import PropertyColumnMapping 

34from sonar_backend.settings import KEEP_IMPORTED_DATA_FILES 

35from sonar_backend.settings import LOGGER 

36from sonar_backend.settings import PROFILE_IMPORT 

37from sonar_backend.settings import PROPERTY_BATCH_SIZE 

38from sonar_backend.settings import REDIS_URL 

39from sonar_backend.settings import SAMPLE_BATCH_SIZE 

40from sonar_backend.settings import SONAR_DATA_ARCHIVE 

41from sonar_backend.settings import SONAR_DATA_ENTRY_FOLDER 

42from sonar_backend.settings import SONAR_DATA_PROCESSING_FOLDER 

43 

44property_cache = {} 

45 

46 

47def check_for_new_data(): 

48 """Checks for new processing jobs in the database queue and processes them sequentially. 

49 Steps: 

50 - Fetches jobs with `QUEUED` status from the 'ProcessingJob' table, ordered by 'entry_time'. 

51 - For each job: 

52 - Retrieves associated files from 'FileProcessing'. 

53 - If an error occurs (files not found), updates job status to 'FAILED'. 

54 - Moves valid files to the processing directory. 

55 - Calls 'import_archive' to handle import logic. 

56 - Recursively calls itself to process subsequent jobs. 

57 

58 """ 

59 processing_dir = pathlib.Path(SONAR_DATA_PROCESSING_FOLDER) 

60 if REDIS_URL is None: 

61 print("--------------- WARNING -----------------") 

62 print("REDIS_URL not set, running without celery.") 

63 print("This will take a long time.") 

64 print("--------------- ------- -----------------") 

65 # only one worker (Master) can take a job at a time 

66 with transaction.atomic(): 

67 # Fetch and lock a single job with QUEUED status NOTE: Order of entrytime is matter!! 

68 job = ( 

69 ProcessingJob.objects.select_for_update() 

70 .filter(status=ProcessingJob.ImportType.QUEUED) 

71 .order_by("entry_time") 

72 .first() 

73 ) 

74 if job is None: 

75 return 

76 # Update job status to IN_PROGRESS 

77 job.status = ProcessingJob.ImportType.IN_PROGRESS 

78 job.save() 

79 

80 LOGGER.info(f"## New processing job: {job.job_name} ---") 

81 files = FileProcessing.objects.filter(processing_job=job) 

82 if not files.exists(): 

83 LOGGER.warning( 

84 f"No associated files found, in {job.job_name}, marking as FAILED, continue the proces.." 

85 ) 

86 job.status = ProcessingJob.ImportType.FAILED 

87 job.save() 

88 return 

89 

90 for file in files: 

91 file_path = pathlib.Path(SONAR_DATA_ENTRY_FOLDER).joinpath(file.file_name) 

92 if not file_path.exists(): 

93 LOGGER.error(f"File {file_path} does not exist, marking job as FAILED!") 

94 job.status = ProcessingJob.ImportType.FAILED 

95 job.save() 

96 return 

97 

98 # move files to SONAR_DATA_PROCESSING_FOLDER 

99 if "_prop" in str(file_path): 

100 # Handle corresponding .pkl file for property imports 

101 pkl_file = file_path.with_suffix(".pkl") 

102 new_zip_path = file_path.rename(processing_dir.joinpath(file_path.name)) 

103 new_pkl_path = pkl_file.rename(processing_dir.joinpath(pkl_file.name)) 

104 import_archive( 

105 new_zip_path, pkl_path=new_pkl_path 

106 ) # Pass the .pkl path to import_archive 

107 else: 

108 # process as normal sample or annotation import 

109 new_zip_path = file_path.rename(processing_dir.joinpath(file_path.name)) 

110 import_archive(new_zip_path) 

111 

112 # Recursively check for new data after processing current batch 

113 check_for_new_data() 

114 print("---- No new jobs found. ----") 

115 

116 

117def import_archive(process_file_path: pathlib.Path, pkl_path: pathlib.Path = None): 

118 """Processes an archive file by extracting and importing its contents. 

119 Steps: 

120 - Determines the corresponding `ProcessingJob` based on the file. 

121 - Updates the job status to `IN_PROGRESS`. 

122 - Extracts the archive into a temporary directory. 

123 - Differentiates between property imports (requires `.pkl` mapping) and sample/annotation imports. 

124 - Uses Celery (if available) to process batches; otherwise, processes sequentially. 

125 - On success, moves files to `completed` directory or deletes them. 

126 - On failure, moves files to `error` directory and logs the error. 

127 - Updates the job status based on the completion of the processed file. 

128 

129 NOTE: [optional] if it just starts, we update the ProcessingJob (to IP) 

130 """ 

131 temp_dir = ( 

132 pathlib.Path(SONAR_DATA_ARCHIVE) 

133 .joinpath("temp") 

134 .joinpath(process_file_path.stem) 

135 ) 

136 try: 

137 filename_ID = process_file_path.name 

138 # get JOB ID based on the given files 

139 proJob_obj = ProcessingJob.objects.filter(files__file_name=filename_ID).first() 

140 if not proJob_obj: 

141 LOGGER.warning( 

142 "The given import files is not related to any jobID (skip this batch)" 

143 ) 

144 return 

145 job_ID = proJob_obj.job_name 

146 LOGGER.info(f"Process job: {job_ID}") 

147 print(f"Unzip to {temp_dir}") 

148 # ProcessingJob.objects.filter(job_name=job_ID).update( 

149 # status=ProcessingJob.ImportType.IN_PROGRESS 

150 # ) 

151 # unzip the zip file to SONAR_DATA_ARCHIVE 

152 with zipfile.ZipFile(process_file_path, "r") as zip_ref: 

153 zip_ref.extractall(temp_dir) 

154 

155 # Files are distributed and processed 

156 print(f"Running data entry for {temp_dir}") 

157 if pkl_path and pkl_path.exists(): 

158 # property import 

159 import_type = ImportLog.ImportType.PROPERTY 

160 print(f"Property import detected") 

161 batch_size = PROPERTY_BATCH_SIZE 

162 print("Batch size:", batch_size) 

163 property_files_tsv = list(temp_dir.glob("**/*.tsv")) 

164 property_files_csv = list(temp_dir.glob("**/*.csv")) 

165 property_files = ( 

166 property_files_tsv + property_files_csv 

167 ) # Combine both types 

168 

169 if property_files: 

170 # Load the column mapping from the .pkl file 

171 with open(pkl_path, "rb") as pkl_file: 

172 column_mapping = pickle.load( 

173 pkl_file 

174 ) # Load the .pkl file into column_mapping 

175 

176 use_celery = bool(REDIS_URL) # Use Celery if Redis is configured 

177 for property_file in property_files: 

178 sep = "," if property_file.suffix == ".csv" else "\t" 

179 import_property( 

180 str(property_file), 

181 sep, 

182 use_celery, 

183 column_mapping, 

184 batch_size=batch_size, 

185 ) # Pass column_mapping 

186 

187 else: 

188 batch_size = SAMPLE_BATCH_SIZE 

189 print("Batch size:", batch_size) 

190 # var and vcf 

191 sample_files = list(temp_dir.joinpath("samples").glob("**/*.sample")) 

192 anno_files = list(temp_dir.joinpath("anno").glob("**/*.vcf.*")) 

193 print(f"Sample: {len(sample_files)} files found") 

194 print(f"Annotation (vcfs): {len(anno_files)} files found") 

195 if len(sample_files) > 0: 

196 import_type = ImportLog.ImportType.SAMPLE 

197 elif len(anno_files) > 0: 

198 import_type = ImportLog.ImportType.ANNOTATION 

199 else: 

200 import_type = ImportLog.ImportType.SAMPLE_ANNOTATION_ARCHIVE 

201 

202 timer = datetime.now() 

203 

204 number_of_batches = ( 

205 (len(sample_files) + batch_size - 1) // batch_size 

206 if sample_files 

207 else (len(anno_files) + batch_size - 1) // batch_size 

208 ) 

209 

210 print(f"Total number of batches: {number_of_batches}") 

211 sample_files = [str(file) for file in sample_files] 

212 if batch_size: 

213 replicon_cache = {} 

214 gene_cache_by_accession = {} 

215 gene_cache_by_var_pos = {} 

216 if REDIS_URL: 

217 print("setting up sample import celery jobs..") 

218 sample_jobs = [] 

219 for i in range(0, len(sample_files), batch_size): 

220 batch = sample_files[i : i + batch_size] 

221 sample_jobs.append( 

222 process_batch.s( 

223 batch, 

224 replicon_cache, 

225 gene_cache_by_accession, 

226 gene_cache_by_var_pos, 

227 str(temp_dir), 

228 ) 

229 ) 

230 results = group(sample_jobs).apply_async().get() 

231 for result in results: 

232 if not result[0]: 

233 raise Exception( 

234 f"Sample Import Error: {result[1]} - {result[2]}" 

235 ) 

236 results = ( 

237 group([process_annotation.s(str(file)) for file in anno_files]) 

238 .apply_async() 

239 .get() 

240 ) 

241 for result in results: 

242 if not result[0]: 

243 raise Exception( 

244 f"Annotation Import Error: {result[1]} - {result[2]}" 

245 ) 

246 else: 

247 replicon_cache = {} 

248 gene_cache_by_accession = {} 

249 # Samples 

250 for i in range(0, len(sample_files), batch_size): 

251 # print( 

252 # f"processing batch {(i//batch_size) + 1} of {number_of_batches}" 

253 # ) 

254 # batchtimer = datetime.now() 

255 batch = sample_files[i : i + batch_size] 

256 process_batch_single_thread( 

257 batch, 

258 replicon_cache, 

259 gene_cache_by_accession, 

260 gene_cache_by_var_pos, 

261 str(temp_dir), 

262 ) 

263 # annotation 

264 for file in anno_files: 

265 process_annotation(str(file)) 

266 

267 # print( 

268 # f"batch {(i//batch_size) + 1} done in {datetime.now() - batchtimer}" 

269 # ) 

270 

271 LOGGER.info(f"import done in {datetime.now() - timer}") 

272 except Exception as e: 

273 # TODO: [optional] if fail, we update the ProcessingJob right now. 

274 LOGGER.error(f"Error : {e}") 

275 error_dir = pathlib.Path(SONAR_DATA_ARCHIVE).joinpath("error") 

276 error_dir.mkdir(parents=True, exist_ok=True) 

277 process_file_path.rename(error_dir.joinpath(process_file_path.name)) 

278 if pkl_path and pkl_path.exists(): 

279 pkl_path.rename(error_dir.joinpath(pkl_path.name)) 

280 ImportLog.objects.create( 

281 type=import_type, 

282 file=FileProcessing.objects.get(file_name=filename_ID), 

283 success=False, 

284 exception_text=e, 

285 stack_trace=traceback.format_exc(), 

286 ) 

287 LOGGER.error(f"--- Exception: move to {error_dir} ---") 

288 

289 else: # no exception occurs 

290 if KEEP_IMPORTED_DATA_FILES: 

291 completed_dir = pathlib.Path(SONAR_DATA_ARCHIVE).joinpath("completed") 

292 completed_dir.mkdir(parents=True, exist_ok=True) 

293 process_file_path.rename(completed_dir.joinpath(process_file_path.name)) 

294 if pkl_path and pkl_path.exists(): 

295 pkl_path.rename(completed_dir.joinpath(pkl_path.name)) 

296 LOGGER.info(f"--- Finish: move to {completed_dir} ---") 

297 else: 

298 # Delete files after import rather than keeping them in the 

299 # archive/completed/ dir 

300 process_file_path.unlink() 

301 if pkl_path and pkl_path.exists(): 

302 pkl_path.unlink() 

303 ImportLog.objects.create( 

304 type=import_type, 

305 file=FileProcessing.objects.get(file_name=filename_ID), 

306 success=True, 

307 ) 

308 finally: 

309 # TODO: need to rethink about how to finalize the job status 

310 # The _file_count = total_all_file - total_file is not a great idea 

311 shutil.rmtree(temp_dir) 

312 # --- update job status 

313 

314 # list all related files 

315 all_file_obj = FileProcessing.objects.filter(processing_job__job_name=job_ID) 

316 total_all_file = all_file_obj.count() 

317 all_file_names_list = all_file_obj.values_list("file_name", flat=True) 

318 # Import search only success and count 

319 _file_obj = ImportLog.objects.filter(file__in=all_file_names_list) 

320 total_file = _file_obj.count() 

321 # Calculate the difference between total files and successful imports 

322 _file_count = total_all_file - total_file 

323 # Update ProcessingJob status based on file processing status 

324 if _file_count > 0: 

325 # "In Progress" 

326 ProcessingJob.objects.filter(job_name=job_ID).update( 

327 status=ProcessingJob.ImportType.IN_PROGRESS 

328 ) 

329 # print("In progress: %s", _file_count) 

330 else: 

331 # If all files are successfully processed, check if any errors occurred during import 

332 error_files = ImportLog.objects.filter(file__in=all_file_obj, success=False) 

333 if error_files.exists(): 

334 # "Failed" 

335 ProcessingJob.objects.filter(job_name=job_ID).update( 

336 status=ProcessingJob.ImportType.FAILED 

337 ) 

338 # print(f"### Job ID: {job_ID} \nRun status: Failed") 

339 LOGGER.error(f"### Job ID: {job_ID} \nRun status: Failed") 

340 else: 

341 # "Completed" 

342 ProcessingJob.objects.filter(job_name=job_ID).update( 

343 status=ProcessingJob.ImportType.COMPLETED 

344 ) 

345 LOGGER.info(f"### Job ID: {job_ID} \nRun status: Completed") 

346 

347 

348@shared_task 

349def process_batch( 

350 batch: list[str], 

351 replicon_cache, 

352 gene_cache_by_accession, 

353 gene_cache_by_var_pos, 

354 temp_dir, 

355): 

356 parameters = locals().copy() 

357 if PROFILE_IMPORT: 

358 lp = LineProfiler() 

359 # Add a few of the slowest functions based on profiling the 

360 # process_batch_run() function 

361 lp.add_function(SonarImport.get_mutation_objs_cds_and_parent_relations) 

362 lp.add_function(SonarImport.get_mutation_objs_nt) 

363 lp.add_function(process_batch_run) 

364 lp.add_function(process_annotation) 

365 process_batch_profiled = lp(process_batch_run) 

366 retval = process_batch_profiled(**parameters) 

367 lp.print_stats() 

368 return retval 

369 else: 

370 return process_batch_run(**parameters) 

371 

372 

373def process_batch_run( 

374 batch: list[str], 

375 replicon_cache, 

376 gene_cache_by_accession, 

377 gene_cache_by_var_pos, 

378 temp_dir, 

379): 

380 try: 

381 sonar_import_objs = [ 

382 SonarImport(pathlib.Path(file), import_folder=temp_dir) for file in batch 

383 ] 

384 sequences = [ 

385 sample_import_obj.get_sequence_obj() 

386 for sample_import_obj in sonar_import_objs 

387 ] 

388 

389 # Use bulk upsert 

390 # performs INSERT ... ON CONFLICT DO UPDATE 

391 with cache.lock("sequence"): 

392 Sequence.objects.bulk_create( 

393 sequences, 

394 update_conflicts=True, 

395 unique_fields=["name"], 

396 update_fields=["seqhash", "length", "last_update_date"], 

397 ) 

398 alignments: list[Alignment] = [] 

399 for sample_import_obj in sonar_import_objs: 

400 sample_import_obj.update_replicon_obj(replicon_cache) 

401 sample_import_obj.create_alignment(alignments) 

402 

403 with cache.lock("alignment"): 

404 Alignment.objects.bulk_create( 

405 alignments, 

406 update_conflicts=True, 

407 unique_fields=["sequence", "replicon"], 

408 update_fields=["sequence", "replicon"], 

409 ) 

410 

411 nt_mutation_set: list[NucleotideMutation] = [] 

412 cds_mutation_set: list[AminoAcidMutation] = [] 

413 mutation_parent_relations = [] 

414 nt_mutation_alignment_relations: list[NucleotideMutation.alignments.through] = ( 

415 [] 

416 ) 

417 aa_mutation_alignment_relations: list[AminoAcidMutation.alignments.through] = [] 

418 for sample_import_obj in sonar_import_objs: 

419 id_to_mutation_mapping = sample_import_obj.get_mutation_objs_nt( 

420 nt_mutation_set, 

421 replicon_cache, 

422 gene_cache_by_var_pos, 

423 nt_mutation_alignment_relations, 

424 ) 

425 parent_relations = ( 

426 sample_import_obj.get_mutation_objs_cds_and_parent_relations( 

427 cds_mutation_set, 

428 gene_cache_by_accession, 

429 id_to_mutation_mapping, 

430 aa_mutation_alignment_relations, 

431 ) 

432 ) 

433 mutation_parent_relations.extend(parent_relations) 

434 with cache.lock("mutation"): 

435 NucleotideMutation.objects.bulk_create( 

436 nt_mutation_set, 

437 update_conflicts=True, 

438 unique_fields=["ref", "alt", "start", "end", "replicon"], 

439 update_fields=[ 

440 "is_frameshift", 

441 ], 

442 ) 

443 

444 AminoAcidMutation.objects.bulk_create( 

445 cds_mutation_set, 

446 update_conflicts=True, 

447 unique_fields=["ref", "alt", "start", "end", "cds"], 

448 update_fields=["ref", "alt", "start", "end", "cds"], 

449 ) 

450 AminoAcidMutation.parent.through.objects.bulk_create( 

451 mutation_parent_relations, 

452 ignore_conflicts=True, 

453 ) 

454 NucleotideMutation.alignments.through.objects.bulk_create( 

455 [ 

456 NucleotideMutation.alignments.through( 

457 nucleotidemutation_id=rel.nucleotidemutation.id, 

458 alignment_id=rel.alignment.id, 

459 ) 

460 for rel in nt_mutation_alignment_relations 

461 ], 

462 ignore_conflicts=True, 

463 ) 

464 AminoAcidMutation.alignments.through.objects.bulk_create( 

465 [ 

466 AminoAcidMutation.alignments.through( 

467 aminoacidmutation_id=rel.aminoacidmutation.id, 

468 alignment_id=rel.alignment.id, 

469 ) 

470 for rel in aa_mutation_alignment_relations 

471 ], 

472 ignore_conflicts=True, 

473 ) 

474 

475 return (True, None, None) 

476 

477 except Exception as e: 

478 # Handle the DataError exception here 

479 LOGGER.error(f"DataError: {e}") 

480 # Perform additional error handling or logging as needed 

481 LOGGER.error("Error happens on this batch") 

482 return (False, str(e), traceback.format_exc()) 

483 

484 

485@shared_task 

486def process_annotation(file_name): 

487 try: 

488 annotation_import = AnnotationImport(file_name) 

489 if REDIS_URL: 

490 with cache.lock("annotation"): 

491 AnnotationType.objects.bulk_create( 

492 annotation_import.get_annotation_objs(), 

493 ignore_conflicts=True, 

494 ) 

495 with cache.lock("annotation2mutation"): 

496 AnnotationType.mutations.through.objects.bulk_create( 

497 annotation_import.get_annotation2mutation_objs(), 

498 ignore_conflicts=True, 

499 ) 

500 else: 

501 AnnotationType.objects.bulk_create( 

502 annotation_import.get_annotation_objs(), 

503 ignore_conflicts=True, 

504 ) 

505 AnnotationType.mutations.through.objects.bulk_create( 

506 annotation_import.get_annotation2mutation_objs(), ignore_conflicts=True 

507 ) 

508 except Exception as e: 

509 LOGGER.error(f"Error in process_annotation: {e}") 

510 return (False, str(e), traceback.format_exc()) 

511 return (True, None, None) 

512 

513 

514def process_batch_single_thread( 

515 batch, replicon_cache, gene_cache_by_accession, gene_cache_by_var_pos, temp_dir 

516): 

517 try: 

518 sample_import_objs = [ 

519 SonarImport(file, import_folder=temp_dir) for file in batch 

520 ] 

521 with transaction.atomic(): 

522 sequences = [ 

523 sample_import_obj.get_sequence_obj() 

524 for sample_import_obj in sample_import_objs 

525 ] 

526 Sequence.objects.bulk_create( 

527 sequences, 

528 update_conflicts=True, 

529 unique_fields=["name"], # Use name as the unique identifier 

530 update_fields=[ 

531 "seqhash", 

532 "length", 

533 "last_update_date", 

534 ], # Update these fields 

535 ) 

536 [x.update_replicon_obj(replicon_cache) for x in sample_import_objs] 

537 alignments: list[Alignment] = [] 

538 for sample_import_obj in sample_import_objs: 

539 sample_import_obj.create_alignment(alignments) 

540 Alignment.objects.bulk_create( 

541 alignments, 

542 update_conflicts=True, 

543 unique_fields=["sequence", "replicon"], 

544 update_fields=["sequence", "replicon"], 

545 ) 

546 nt_mutation_set: list[NucleotideMutation] = [] 

547 cds_mutation_set: list[AminoAcidMutation] = [] 

548 mutation_parent_relations = [] 

549 nt_mutation_alignment_relations: list[ 

550 NucleotideMutation.alignments.through 

551 ] = [] 

552 aa_mutation_alignment_relations: list[ 

553 AminoAcidMutation.alignments.through 

554 ] = [] 

555 

556 for sample_import_obj in sample_import_objs: 

557 id_to_mutation_mapping = sample_import_obj.get_mutation_objs_nt( 

558 nt_mutation_set, 

559 replicon_cache, 

560 gene_cache_by_var_pos, 

561 nt_mutation_alignment_relations, 

562 ) 

563 parent_relations = ( 

564 sample_import_obj.get_mutation_objs_cds_and_parent_relations( 

565 cds_mutation_set, 

566 gene_cache_by_accession, 

567 id_to_mutation_mapping, 

568 aa_mutation_alignment_relations, 

569 ) 

570 ) 

571 mutation_parent_relations.extend(parent_relations) 

572 

573 NucleotideMutation.objects.bulk_create( 

574 nt_mutation_set, 

575 update_conflicts=True, 

576 unique_fields=["ref", "alt", "start", "end", "replicon"], 

577 update_fields=["ref", "alt", "start", "end", "replicon"], 

578 ) 

579 AminoAcidMutation.objects.bulk_create( 

580 cds_mutation_set, 

581 update_conflicts=True, 

582 unique_fields=["ref", "alt", "start", "end", "cds"], 

583 update_fields=["ref", "alt", "start", "end", "cds"], 

584 ) 

585 AminoAcidMutation.parent.through.objects.bulk_create( 

586 mutation_parent_relations, 

587 ignore_conflicts=True, 

588 ) 

589 NucleotideMutation.alignments.through.objects.bulk_create( 

590 [ 

591 NucleotideMutation.alignments.through( 

592 nucleotidemutation_id=rel.nucleotidemutation.id, 

593 alignment_id=rel.alignment.id, 

594 ) 

595 for rel in nt_mutation_alignment_relations 

596 ], 

597 ignore_conflicts=True, 

598 ) 

599 AminoAcidMutation.alignments.through.objects.bulk_create( 

600 [ 

601 AminoAcidMutation.alignments.through( 

602 aminoacidmutation_id=rel.aminoacidmutation.id, 

603 alignment_id=rel.alignment.id, 

604 ) 

605 for rel in aa_mutation_alignment_relations 

606 ], 

607 ignore_conflicts=True, 

608 ) 

609 

610 # annotations = [] 

611 # for sample_import_obj in sample_import_objs: 

612 # annotations.extend(sample_import_obj.get_annotation_objs()) 

613 # AnnotationType.objects.bulk_create( 

614 # annotations, 

615 # ignore_conflicts=True, 

616 # ) 

617 # annotation2mutations = [] 

618 # for sample_import_obj in sample_import_objs: 

619 # annotation2mutations.extend( 

620 # sample_import_obj.get_annotation2mutation_objs() 

621 # ) 

622 # Mutation2Annotation.objects.bulk_create( 

623 # annotation2mutations, ignore_conflicts=True 

624 # ) 

625 

626 # Filter sequences without associated samples 

627 # clean_unused_sequences() 

628 

629 except DataError as data_error: 

630 # Handle the DataError exception here 

631 LOGGER.error(f"DataError: {data_error}") 

632 # Perform additional error handling or logging as needed 

633 LOGGER.error("Error happens on this batch") 

634 for sample_import_obj in sample_import_objs: 

635 LOGGER.error(f"{sample_import_obj.sample_file_path}") 

636 raise 

637 except Exception as e: 

638 # Handle other exceptions if necessary 

639 LOGGER.critical(f"An unexpected error occurred: {e}") 

640 raise 

641 

642 

643def import_property( 

644 property_file, sep, use_celery=False, column_mapping=None, batch_size=1000 

645): 

646 try: 

647 # Load the CSV file in batches 

648 properties_df = pd.read_csv( 

649 property_file, 

650 sep=sep, 

651 dtype="string", 

652 keep_default_na=False, 

653 ).replace("", pd.NA) 

654 timer = datetime.now() 

655 

656 # Use Celery if parallel processing is enabled 

657 # Convert column_mapping object (PropertyColumnMapping class) to a JSON-serializable format 

658 # because celery takes JSON-serializable data (like dictionaries or lists) and passes to the Celery tasks 

659 if column_mapping is not None: 

660 serializable_column_mapping = { 

661 k: { 

662 "db_property_name": v.db_property_name, 

663 "data_type": v.data_type, 

664 "default": v.default, 

665 } 

666 for k, v in column_mapping["column_mapping"].items() 

667 } 

668 sample_name_column = column_mapping["sample_id_column"] 

669 related_sequences_column = column_mapping["sequences_id_column"] 

670 else: 

671 serializable_column_mapping = None 

672 sample_name_column = "ID" 

673 related_sequences_column = "ID" 

674 

675 # Data preprocessing 

676 

677 for column_name, col_info in serializable_column_mapping.items(): 

678 if column_name not in properties_df.columns: 

679 print( 

680 f"Skipping column '{column_name}' as it is not in the property file." 

681 ) 

682 continue 

683 # Apply default values for missing or empty fields 

684 default_value = col_info["default"] 

685 print(f"{column_name} :pairs with: {col_info}") 

686 

687 if default_value is not None: 

688 properties_df[column_name] = properties_df[column_name].fillna( 

689 default_value 

690 ) # Replace NaN values 

691 

692 # # Explicitly convert the column to string to avoid .0 issues 

693 # if ( 

694 # col_info["data_type"] == "value_varchar" 

695 # and column_name in properties_df.columns 

696 # ): 

697 # properties_df[column_name] = properties_df[column_name].astype(str) 

698 

699 # Format columns with 'value_date' type 

700 if ( 

701 col_info["data_type"] == "value_date" 

702 and column_name in properties_df.columns 

703 ): 

704 properties_df[column_name] = properties_df[column_name].apply( 

705 parse_date 

706 ) 

707 if use_celery: 

708 print("Setting up property import celery jobs...") 

709 

710 property_jobs = [] 

711 for i in range(0, len(properties_df), batch_size): 

712 batch = properties_df.iloc[i : i + batch_size] 

713 if isinstance(batch, pd.Series): 

714 batch = ( 

715 batch.to_frame().T 

716 ) # Ensure it's a DataFrame if it's a Series 

717 batch_as_dict = batch.to_dict(orient="records") 

718 property_jobs.append( 

719 process_property_batch.s( 

720 batch_as_dict, 

721 sample_name_column, 

722 serializable_column_mapping, 

723 related_sequences_column, 

724 ) # Pass column_mapping 

725 ) 

726 

727 # Run the Celery group of tasks and collect results 

728 results = group(property_jobs).apply_async().get() 

729 

730 for result in results: 

731 if not result[0]: 

732 raise Exception(f"Property Import Error: {result[1]}") 

733 

734 else: 

735 print("Processing properties in single-threaded mode...") 

736 

737 # Process the CSV file in a single-threaded manner 

738 if isinstance(properties_df, pd.Series): 

739 properties_df = properties_df.to_frame().T 

740 batch_as_dict = properties_df.to_dict(orient="records") 

741 # column_mapping does not need to be JSON-serializable format, because we use only single thread 

742 _process_property_file( 

743 batch_as_dict, 

744 sample_name_column, 

745 column_mapping["column_mapping"], 

746 related_sequences_column, 

747 ) 

748 

749 print(f"Property import usage time: {datetime.now() - timer}") 

750 

751 except Exception as e: 

752 print("Error in import_property func.:", e) 

753 print(f"Error in import_property line#: {e.__traceback__.tb_lineno}") 

754 raise # Re-raise the exception to propagate it up the stack? 

755 

756 

757@shared_task 

758def process_property_batch( 

759 batch_as_dict, 

760 sample_name_column, 

761 serialized_column_mapping, 

762 related_sequences_column, 

763): 

764 # Reconstruct column_mapping from the serialized format 

765 if serialized_column_mapping is not None: 

766 column_mapping = { 

767 k: PropertyColumnMapping( 

768 db_property_name=v["db_property_name"], 

769 data_type=v["data_type"], 

770 default=v["default"], 

771 ) 

772 for k, v in serialized_column_mapping.items() 

773 } 

774 else: 

775 column_mapping = None 

776 

777 try: 

778 _process_property_file( 

779 batch_as_dict, sample_name_column, column_mapping, related_sequences_column 

780 ) 

781 except Exception as e: 

782 print("Error in process_property_batch func.:", e) 

783 print(f"Error in process_property_batch line#: {e.__traceback__.tb_lineno}") 

784 return False, f"Failed or interrupted batch job: {e}" 

785 return True, "Batch processed successfully" 

786 

787 

788def _process_property_file( 

789 batch_as_dict, sample_name_column, column_mapping, related_sequences_column 

790): 

791 """ 

792 Logic for processing a batch of the property file 

793 """ 

794 properties_df = pd.DataFrame.from_dict(batch_as_dict, dtype=object) 

795 properties_df.drop_duplicates( 

796 subset=[sample_name_column], keep="last", inplace=True 

797 ) 

798 if sample_name_column == related_sequences_column: 

799 properties_df[f"{sample_name_column}_name"] = properties_df[sample_name_column] 

800 sample_name_column = f"{sample_name_column}_name" 

801 properties_df[related_sequences_column] = properties_df[ 

802 related_sequences_column 

803 ].apply(lambda x: x.split() if x else []) 

804 sample_property_names = [] 

805 custom_property_names = [] 

806 

807 # Separate sample properties and custom properties 

808 for property_name in properties_df.columns: 

809 if property_name in column_mapping.keys(): 

810 db_property_name = column_mapping[property_name].db_property_name 

811 try: 

812 models.Sample._meta.get_field(db_property_name) 

813 sample_property_names.append(db_property_name) 

814 except FieldDoesNotExist: 

815 custom_property_names.append(property_name) 

816 sample_name_set = set(properties_df[sample_name_column]) 

817 properties_df.convert_dtypes() 

818 properties_df.set_index(sample_name_column, inplace=True) 

819 

820 # Update existing samples 

821 # Fetch existing samples from the database (empty before first property import) 

822 samples = models.Sample.objects.filter(name__in=sample_name_set).iterator() 

823 sample_updates = [] 

824 property_updates = [] 

825 existing_sample_names = set() 

826 try: 

827 for sample in samples: 

828 existing_sample_names.add(sample.name) 

829 row = properties_df.loc[sample.name] 

830 sample.last_update_date = timezone.now() 

831 for name, value in row.items(): 

832 if name in column_mapping.keys(): 

833 db_name = column_mapping[name].db_property_name 

834 if db_name in sample_property_names: 

835 setattr(sample, db_name, value) 

836 sequence_ids = row[related_sequences_column] 

837 sequences = models.Sequence.objects.filter(name__in=sequence_ids)