diff --git a/management/commands/indexer.py b/management/commands/indexer.py index 5f260ae..73afa2c 100644 --- a/management/commands/indexer.py +++ b/management/commands/indexer.py @@ -3,6 +3,7 @@ from GaiaDBInterface.models import GaiaSource, CatalogFile from astropy.coordinates import SkyCoord import astropy_healpix as ah import numpy as np +from datetime import datetime, timedelta def current_time(): return (datetime.now() + timedelta(hours=3)).strftime("%H:%M:%S") @@ -13,7 +14,7 @@ def healpix(): ORDER='ring' for catalog_file in CatalogFile.objects.all(): - + print(f'[{current_time()}] Loading sources from the database...') sources = list(catalog_file.sources.all()) print(f'[{current_time()}] Sources ready. Forming ra & dec arrays...') diff --git a/management/commands/ingester.py b/management/commands/ingester.py index 6bcde2f..9065df1 100644 --- a/management/commands/ingester.py +++ b/management/commands/ingester.py @@ -6,6 +6,7 @@ import json import glob import uuid import asyncio +import select import tracemalloc from datetime import datetime, timedelta import pandas as pd @@ -14,7 +15,13 @@ from asgiref.sync import sync_to_async from django.core.management.base import BaseCommand from GaiaDBInterface.models import GaiaSource, CatalogFile - +def input_with_timeout(prompt, timeout=30): + print(prompt, end='', flush=True) + ready, _, _ = select.select([sys.stdin], [], [], timeout) + if ready: + return sys.stdin.readline().strip() + else: + return '' class Command(BaseCommand): help = 'Ingest CSV files into the database' @@ -32,7 +39,7 @@ class Command(BaseCommand): #fetching the file list - directory = input(f"Please enter the path to the directory containing the csv files [{previous_choice}]: ") + directory = input_with_timeout(f"Please enter the path to the directory containing the csv files [{previous_choice}]: ", 30) if not directory: directory = previous_choice @@ -66,7 +73,7 @@ class Command(BaseCommand): - input("Press Enter to continue...") + #input("Press Enter to continue...") @@ -157,6 +164,10 @@ class Command(BaseCommand): current_time = (datetime.now() + timedelta(hours=3)).strftime("%H:%M:%S") #Timestamp and progress self.stdout.write statement self.stdout.write(f"[{current_time}] {ingested_files_count}/{len(catalog_files)}") - + + # Create a completion flag file + with open("ingester_done.flag", "w") as f: + f.write("done") + print("Ingester done") asyncio.run(ingest_files())