Update SMP scripts

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)

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()
 

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()

snowex_db/metadata_variable_overrides.yaml

@@ -15,3 +15,4 @@ IGNORE:
     - timing
     - original_total_samples
     - data_subsampled_to
+    - total_samples

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")

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