Skip to content

Commit

Permalink
Allow sasquatch datastore to add extra fields
Browse files Browse the repository at this point in the history
  • Loading branch information
timj committed Mar 21, 2024
1 parent 11c21f9 commit 9eaf31e
Showing 1 changed file with 21 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,16 @@ class SasquatchDatastore(GenericBaseDatastore):
environment variable and "lsst.dm" if no environment variable is defined.
"""

extra_fields: dict[str, str | int | float] | None
"""Extra key/value pairs that should be passed along with the metric
when storing in Sasquatch.
If not given extra key/value pairs can be obtained from the
``$SASQUATCH_EXTRAS`` environment variable. This environment variable
must have the form of "k1=v2;k2=v2". If no environment variable is given
defaults can be read from the datastore configuration.
"""

def __init__(
self,
config: DatastoreConfig,
Expand Down Expand Up @@ -113,6 +123,16 @@ def __init__(
),
)

extra_fields: dict[str, str | int | float] | None = None
if extras_str := os.environ.get("SASQUATCH_EXTRAS"):
extra_fields = {}
for item in extras_str.split(";"):
k, v = item.split("=")
extra_fields[k] = v
self.extra_fields = extra_fields
else:
self.extra_fields = self.config.get("extra_fields", extra_fields)

self._dispatcher = SasquatchDispatcher(self.restProxyUrl, self.accessToken, self.namespace)

@classmethod
Expand All @@ -133,7 +153,7 @@ def bridge(self) -> DatastoreRegistryBridge:

def put(self, inMemoryDataset: Any, ref: DatasetRef) -> None:
if self.constraints.isAcceptable(ref):
self._dispatcher.dispatchRef(inMemoryDataset, ref)
self._dispatcher.dispatchRef(inMemoryDataset, ref, extraFields=self.extra_fields)
else:
log.debug("Could not put dataset type %s with Sasquatch datastore", ref.datasetType)
raise DatasetTypeNotSupportedError(
Expand Down

0 comments on commit 9eaf31e

Please sign in to comment.