Skip to content
Snippets Groups Projects
flag_spikes.py 3.81 KiB
Newer Older
  • Learn to ignore specific revisions
  • import os
    from datetime import datetime
    
    import typer
    
    from typing import Optional, List
    
    
    from ..metadata import Metadata, MetadataFactory
    from ..algorithm.SpikesAlgorithm import SpikesAlgorithm
    
    
    def flag_spikes(
        observatory: str = typer.Option(..., help="Observatory code"),
    
        channels: List[str] = typer.Option(
            default=["F"], help="Channels to despike, default is F."
        ),
    
        starttime: datetime = typer.Option(
            default=None,
            help="Start time for metadata, default is None. CLI example usage: --starttime '2024-01-01' ",
        ),
        endtime: datetime = typer.Option(
            default=None,
            help="End time for metadata, default is None. CLI example usage: --endtime '2024-01-01' ",
        ),
        window_size: Optional[int] = typer.Option(
            default=10,
            help="Window size for despike algorithm",
        ),
        threshold: Optional[int] = typer.Option(
            default=3,
            help="Threshold for spike size",
        ),
        metadata_token: str = typer.Option(
            default=os.getenv("GITLAB_API_TOKEN"),
            help="Token for metadata web service.",
            metavar="TOKEN",
            show_default="environment variable GITLAB_API_TOKEN",
        ),
        metadata_url: str = typer.Option(
            default="https://staging-geomag.cr.usgs.gov/ws/secure/metadata",
            help="URL to metadata web service",
            metavar="URL",
        ),
        force: bool = typer.Option(
            default=False,
            help="Force skip the checks",
        ),
    ):
        """Flag spikes and create spike metadata in the metadata webservice."""
    
        # TODO: Consider adding below variables as options
        # run spike detection algorithm on data and return spike metadata
        spike_algorithm = SpikesAlgorithm(
            starttime=starttime,
            endtime=endtime,
            observatory=observatory,
            window_size=window_size,
            threshold=threshold,
    
            channels=channels,
    
            type="variation",
            interval="second",
            metadata_token=metadata_token,
            metadata_url=metadata_url,
        )
    
        spikes = spike_algorithm.run()
    
        # confirm whether or not to create spike metadata
        if not force:
            typer.confirm(f"Are you sure you want to create flag metadata?", abort=True)
            print("Creating flag metadata")
    
        # write flag metadata to metadata service
        metadata_factory = MetadataFactory(token=metadata_token, url=metadata_url)
        with typer.progressbar(
            iterable=spikes, label="Uploading to metadata service"
        ) as progressbar:
            for spike in progressbar:
                upload_spike_metadata(algorithm=spike_algorithm, spike=spike)
    
    
    def main() -> None:
        """Entrypoint for flag spikes method."""
        # for one command, can use typer.run
        typer.run(flag_spikes)
    
    
    def upload_spike_metadata(algorithm: SpikesAlgorithm, spike: Metadata) -> Metadata:
        """Upload new or existing spike metadata to metadata service.
    
        Parameters
        ----------
        algorithm: SpikesAlgorithm
            algorithm to update/load spike metadata
        spike:  Metadata
            spike metadata to be uploaded.
    
        Returns
        -------
        Metadata
            created metadata object.
        """
    
        # check if metadata already exists for period before uploading
        prior_metadata_ID = SpikesAlgorithm.check_existing_metadata(
            algorithm, metadata_obj=spike
        )
        if prior_metadata_ID:
            # TODO: Confirm whether or not to add force or simply update automatically
            typer.confirm(
                f"Spikes already exist for this period, would you like to update this metadata?",
                abort=True,
            )
    
            # set spike.id to prior_metadata.id to update that metadata object in the database
            spike.id = prior_metadata_ID
            return SpikesAlgorithm.update_metadata(algorithm, spike)
        else:
            return SpikesAlgorithm.create_new_metadata(algorithm, spike)
    
    
    if __name__ == "__main__":
        main()