diff --git a/scripts/upload/add_smp.py b/scripts/upload/add_smp.py old mode 100644 new mode 100755 index a053d7b..d00bcab --- a/scripts/upload/add_smp.py +++ b/scripts/upload/add_smp.py @@ -1,91 +1,74 @@ """ -Added smp measurements to the database - -1. Admin must download the NSIDC package first via sh ../download/download_nsidc.sh -2. Run the resample script at least once -3A. python run.py # To run all together all at once -3B. python add_smp.py # To run individually +Added SMP measurements from: +* 2020 Grand Mesa """ +import fileinput -import glob -from os.path import abspath, join - -from snowex_db.batch import UploadProfileBatch -from snowex_db.utilities import get_logger -import concurrent.futures +from earthaccess_data import get_files +from import_logger import get_logger +from snowexsql.db import db_session_with_credentials +from snowex_db.upload.layers import UploadProfileData -def submit_smp(associated_pits, directory, kwargs): - """ - Function for running smp submission threaded - """ - errors = 0 +LOG = get_logger() - for pit_id in associated_pits: +# Map of DATA SET ID to DOI from NSIDC +# * https://nsidc.org/data/snex20_smp/versions/1 +SMP_DOI = { + "SNEX20_SMP": "10.5067/ZYW6IHFRYDSE", +} - # Grab all SMP profiles with this pit_id - pattern = f'*{pit_id}.CSV' - pit_files = glob.glob(join(directory, 'csv_resampled', pattern)) - - # Instantiate the uploader - b = UploadProfileBatch(pit_files, site_id=f'COGM{pit_id}', **kwargs) - - # Submit to the db - b.push() - errors += len(b.errors) +def add_header_indicator(file: str) -> None: + """ + This is a workaround for all SMP files that don't use a '#' for the last header + row with the column names. This information is always on line number 7 when converted + from the original binary file to a CSV. - return errors + Args: + file: File to edit + """ + with fileinput.FileInput(file, inplace=True) as f: + for line_number, line in enumerate(f): + if line_number == 6 and line.startswith('Depth'): + print(f"# {line}", end='') + else: + print(line, end='') +def main(file_list: list, doi: str) -> None: + LOG.info("Starting SMP Upload") -def main(): - # Obtain a list of Grand mesa smp files - directory = abspath('../download/data/SNOWEX/SNEX20_SMP.001') - all_filenames = glob.glob(join(directory, 'csv_resampled', '*.CSV')) + # SMP has binary .PNT and .CSV files + file_list = [file for file in file_list if str(file).endswith(".CSV")] # Keyword arguments. - kwargs = { - # Uploader kwargs - 'debug': True, - - # Constant metadata - 'site_name': 'Grand Mesa', - 'units': 'Newtons', - 'in_timezone': 'UTC', - 'out_timezone': 'UTC', - 'instrument': 'snowmicropen', + smp_metadata = { + 'campaign_name': 'Grand Mesa', + 'timezone': 'UTC', + 'instrument': 'SnowMicroPen', 'header_sep': ':', - 'doi': 'https://doi.org/10.5067/ZYW6IHFRYDSE', - + 'doi': doi, } - # Get logger - log = get_logger('SMP Upload Script') - - # Form the unique pit ids to loop over - associated_pits = list(set(['_'.join(l.split('_')[-2:]).replace('.CSV', '') for l in all_filenames])) - - # Keep track of errors - errors = 0 - nthreads = 6 - pits_per_threads = len(associated_pits) // nthreads - log.info(f'Assigning {pits_per_threads} pits of smp profiles to {nthreads} threads') - - # Loop over by pit ID so we can assign it to groups of files - - with concurrent.futures.ThreadPoolExecutor() as executor: - futures = [] + with db_session_with_credentials() as (_engine, session): + for file in file_list: + add_header_indicator(file) + # SMP data do not have the site ID in the metadata and only in the filename + metadata = file.split("_") + pit_id = metadata[-2] + measurement = metadata[-3] - for i in range(nthreads): - pits = associated_pits[i * pits_per_threads: (i + 1) * pits_per_threads] - futures.append(executor.submit(submit_smp, pits, directory, kwargs)) + LOG.info(f" Adding site {pit_id} and measurement {measurement[-5:]}") - # Collect the errors - for f in futures: - errors += f.result() + # Now make a unique site ID to create an entry per location measurement + pit_id = f"{metadata[-2]}-{measurement[-5:]}" - # Return the number of errors so run.py can report them - return errors + uploader = UploadProfileData( + session, filename=file, id=pit_id, comments=measurement, **smp_metadata + ) + uploader.submit() if __name__ == '__main__': - main() + for data_set_id, doi in SMP_DOI.items(): + with get_files(data_set_id, doi) as files: + main(files, doi) diff --git a/scripts/upload/add_ssa.py b/scripts/upload/add_ssa.py index a981624..7a3da68 100755 --- a/scripts/upload/add_ssa.py +++ b/scripts/upload/add_ssa.py @@ -29,7 +29,7 @@ def main(file_list, doi): for file in file_list: LOG.info(f"Uploading: {file}") uploader = UploadProfileData( - session, filename=str(file), doi=doi, timezone='MST' + session, filename=file, doi=doi, timezone='MST' ) uploader.submit() diff --git a/scripts/upload/earthaccess_data.py b/scripts/upload/earthaccess_data.py index f11e0d3..7f43f77 100644 --- a/scripts/upload/earthaccess_data.py +++ b/scripts/upload/earthaccess_data.py @@ -20,7 +20,9 @@ def get_files(data_set_id:str, doi:str) ->Generator[List[Path], None, None]: files = [] if source_files.exists() and source_files.is_dir(): - files = list(source_files.glob("*.csv")) + files = [ + file.as_posix() for file in source_files.glob("*.csv", case_sensitive=False) + ] if len(files) == 0: earthaccess.login() diff --git a/snowex_db/metadata_variable_overrides.yaml b/snowex_db/metadata_variable_overrides.yaml index 2c99a38..b84c528 100644 --- a/snowex_db/metadata_variable_overrides.yaml +++ b/snowex_db/metadata_variable_overrides.yaml @@ -15,3 +15,4 @@ IGNORE: - timing - original_total_samples - data_subsampled_to + - total_samples diff --git a/snowex_db/upload/base.py b/snowex_db/upload/base.py index f9e16d1..fc41682 100644 --- a/snowex_db/upload/base.py +++ b/snowex_db/upload/base.py @@ -1,6 +1,17 @@ +from dataclasses import dataclass + +@dataclass +class CacheObject: + key: str + id: int + + class BaseUpload: - @staticmethod - def _check_or_add_object(session, clz, check_kwargs, object_kwargs=None): + def __init__(self): + # Lookup cache for inserting + self._lookup_cache = {} + + def _check_or_add_object(self, session, clz, check_kwargs, object_kwargs=None): """ Check for an existing object, add to the database if not found @@ -11,17 +22,26 @@ def _check_or_add_object(session, clz, check_kwargs, object_kwargs=None): object_kwargs: kwargs for instantiating the object """ - # Check if the object exists + # Check in lookup cache + obj = self._lookup_cache.get(str(check_kwargs), None) + if obj: + return obj + + # Check in the database obj = session.query(clz).filter_by(**check_kwargs).first() + + # Create the object or put it in the cache if not obj: # Use check kwargs if not object_kwargs given object_kwargs = object_kwargs or check_kwargs - # If the object does not exist, create it obj = clz(**object_kwargs) session.add(obj) session.commit() + self._lookup_cache[str(check_kwargs)] = CacheObject( + key=str(check_kwargs), id=obj.id + ) return obj @classmethod def _add_entry(cls, **kwargs): - raise NotImplemented("You need this") + raise NotImplementedError("You need this") diff --git a/snowex_db/upload/layers.py b/snowex_db/upload/layers.py index 74798fe..e389803 100644 --- a/snowex_db/upload/layers.py +++ b/snowex_db/upload/layers.py @@ -35,12 +35,10 @@ class UploadProfileData(BaseUpload): expected_attributes = [c for c in dir(LayerData) if c[0] != '_'] TABLE_CLASS = LayerData + INSERT_BATCH_SIZE = 10_000 + def __init__( - self, - session, - filename: Union[str, Path], - timezone: str="US/Mountain", - **kwargs + self, session, filename: Union[str, Path], timezone: str = "US/Mountain", **kwargs ): """ Arguments: @@ -59,6 +57,7 @@ def __init__( instrument_model (str): Instrument name. comments (str): Additional comments. """ + super().__init__() self.log = get_logger(__name__) self.filename = filename @@ -181,17 +180,25 @@ def submit(self): if 'instrument' not in df.columns.values: instrument = self._add_instrument(profile.metadata) - for row in df.to_dict(orient="records"): - if row.get('value') == 'None': + # Skip empty records + df_filtered = df[df['value'] != 'None'] + + all_records_map = [ + self._add_entry(row, campaign, observer_list, site, instrument) + for row in df_filtered.to_dict(orient="records") + ] + + # Process records in batches + for i in range(0, len(all_records_map ), self.INSERT_BATCH_SIZE): + batch = all_records_map [i:i + self.INSERT_BATCH_SIZE] + if not batch: continue - d = self._add_entry( - row, campaign, observer_list, site, instrument, - ) - # session.bulk_save_objects(objects) does not resolve - # foreign keys, DO NOT USE IT - self._session.add(d) - self._session.commit() + self._session.bulk_insert_mappings(self.TABLE_CLASS, batch) + + self._session.commit() + # Mark all cached objects as expired + self._session.expunge_all() else: # procedure to still upload metadata (sites, etc) self.log.warning( @@ -290,13 +297,13 @@ def _add_instrument(self, metadata: SnowExProfileMetadata): Instrument DB record """ # Give priority to passed information from kwargs - instrumen_name = self._instrument or metadata.instrument + instrument_name = self._instrument or metadata.instrument instrument_model = self._instrument_model or metadata.instrument_model return self._check_or_add_object( self._session, Instrument, - dict(name=instrumen_name, model=instrument_model) + dict(name=instrument_name, model=instrument_model) ) def _add_entry( @@ -338,16 +345,14 @@ def _add_entry( ) ) - # Now that the other objects exist and create the entry. - new_entry = self.TABLE_CLASS( - # Required record information + # Create a dictionary for bulk insert + new_entry = dict( depth=row["depth"], bottom_depth=row.get("bottom_depth"), value=row["value"], - # Linked tables - instrument=instrument, - measurement_type=measurement_obj, - site=site, + instrument_id=instrument.id, + measurement_type_id=measurement_obj.id, + site_id=site.id, ) return new_entry diff --git a/snowex_db/upload/points.py b/snowex_db/upload/points.py index 4b3c234..dc2a54f 100644 --- a/snowex_db/upload/points.py +++ b/snowex_db/upload/points.py @@ -64,6 +64,7 @@ def __init__( row_based_timezone instrument_map """ + super().__init__() self.filename = profile_filename self._session = session @@ -206,18 +207,21 @@ def submit(self): c_observations = self._add_campaign_observation(df) measurement_types = self._add_measurement_types(df) + all_records_map = [] for row in df.to_dict(orient="records"): row["geometry"] = WKTElement( str(f"POINT ({row['longitude']} {row['latitude']})"), srid=4326, ) - d = self._add_entry(row, c_observations, measurement_types) + all_records_map.append( + self._add_entry(row, c_observations, measurement_types) + ) - # session.bulk_save_objects(objects) does not resolve - # foreign keys, DO NOT USE IT - self._session.add(d) - self._session.commit() + self._session.bulk_insert_mappings(self.TABLE_CLASS, all_records_map) + self._session.commit() + # Mark all cached objects as expired + self._session.expunge_all() else: # procedure to still upload metadata (sites, etc) LOG.warning( @@ -328,23 +332,21 @@ def _add_campaign_observation(self, df) -> dict: ) date_obj = self._get_first_check_unique(grouped_df, "date") - object_args = dict( + check_args = dict( date=date_obj, name=measurement_name, - # Link objects - doi=doi, - instrument=instrument, + doi_id=doi.id, + instrument_id=instrument.id, ) observation = self._check_or_add_object( self._session, PointObservation, - object_args, + check_args, object_kwargs=dict( - **object_args, + **check_args, description=description, - # Link objects - campaign=campaign, - observer=observer, + campaign_id=campaign.id, + observers_id=observer.id, ) ) @@ -387,7 +389,7 @@ def _add_measurement_types(self, df) -> dict: return types - def _add_entry(self, row: dict, observations: dict, measurement_types: dict) -> PointData: + def _add_entry(self, row: dict, observations: dict, measurement_types: dict) -> dict: """ Add a single point entry and map with the metadata. @@ -409,12 +411,12 @@ def _add_entry(self, row: dict, observations: dict, measurement_types: dict) -> ) # Now that the other objects exist, create the entry - new_entry = self.TABLE_CLASS( + new_entry = dict( datetime=row["datetime"], elevation=row.get('elevation', None), geom=row['geometry'], - measurement_type=measurement_types[row["type"]], - observation=observation, + measurement_type_id=measurement_types[row["type"]].id, + observation_id=observation.id, value=row["value"], )