Coverage for python / lsst / analysis / tools / interfaces / datastore / _dispatcher.py: 12%

298 statements  

« prev     ^ index     » next       coverage.py v7.13.2, created at 2026-02-17 00:47 -0800

1# This file is part of analysis_tools. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ("SasquatchDispatchPartialFailure", "SasquatchDispatchFailure", "SasquatchDispatcher") 

25 

26import calendar 

27import datetime 

28import json 

29import logging 

30import math 

31import re 

32from collections.abc import Mapping, MutableMapping, Sequence 

33from dataclasses import dataclass 

34from typing import TYPE_CHECKING, Any, cast 

35from uuid import UUID, uuid4 

36 

37import numpy as np 

38import requests 

39 

40from lsst.daf.butler import DatasetRef 

41from lsst.resources import ResourcePath 

42from lsst.utils.packages import getEnvironmentPackages 

43 

44from ...utils import http_client 

45 

46if TYPE_CHECKING: 

47 from .. import MetricMeasurementBundle 

48 

49"""Sasquatch datastore""" 

50 

51log = logging.getLogger(__name__) 

52 

53# Constants associated with SasquatchDispatcher 

54PARTITIONS = 1 

55REPLICATION_FACTOR = 3 

56 

57IDENTIFIER_KEYS = [ 

58 "detector", 

59 "patch", 

60 "skymap", 

61 "visit", 

62 "tract", 

63 "physical_filter", 

64 "instrument", 

65 "band", 

66 "exposure", 

67 "group", 

68 "day_obs", 

69] 

70 

71 

72class SasquatchDispatchPartialFailure(RuntimeError): 

73 """This indicates that a Sasquatch dispatch was partially successful.""" 

74 

75 pass 

76 

77 

78class SasquatchDispatchFailure(RuntimeError): 

79 """This indicates that dispatching a 

80 `~lsst.analysis.tool.interface.MetricMeasurementBundle` failed. 

81 """ 

82 

83 pass 

84 

85 

86def _tag2VersionTime(productStr: str) -> tuple[str, float]: 

87 """Determine versions and dates from the string returned from 

88 getEnvironmentPackages. 

89 

90 The `~lsst.utils.packages.genEnvironmentPackages` function returns the 

91 setup version associated with a product, along with a list of tags that 

92 have been added to it. 

93 

94 This method splits up that return string, and determines the earliest date 

95 associated with the setup package version. 

96 

97 Parameters 

98 ---------- 

99 productStr : `str` 

100 The product string returned from a lookup on the result of a call to 

101 `~lsst.utils.packages.getEnvironmentPackages`. 

102 

103 Returns 

104 ------- 

105 result : `tuple` of `str`, `datetime.datetime` 

106 The first `str` is the version of the package, and the second is the 

107 datetime object associated with that released version. 

108 

109 Raises 

110 ------ 

111 ValueError 

112 Raised if there are no tags which correspond to dates. 

113 """ 

114 times: list[datetime.datetime] = [] 

115 version = productStr.split()[0] 

116 tags: str = re.findall("[(](.*)[)]", productStr)[0] 

117 for tag in tags.split(): 

118 numDots = tag.count(".") 

119 numUnder = tag.count("_") 

120 separator = "_" 

121 if numDots > numUnder: 

122 separator = "." 

123 match tag.split(separator): 

124 # Daily tag branch. 

125 case ("d", year, month, day): 

126 dt = datetime.datetime(year=int(year), month=int(month), day=int(day)) 

127 # Weekly tag branch. 

128 case ("w", year, week): 

129 iyear = int(year) 

130 iweek = int(week) 

131 # Use 4 as the day because releases are available starting 

132 # on Thursday 

133 dayOfWeek = 4 

134 

135 # Find the first week to contain a thursday in it 

136 cal = calendar.Calendar() 

137 cal.setfirstweekday(6) 

138 i = 0 

139 for i, iterWeek in enumerate(cal.monthdatescalendar(iyear, 1)): 

140 if iterWeek[dayOfWeek].month == 1: 

141 break 

142 # Handle fromisocalendar not being able to handle week 53 

143 # in the case were the date was going to subtract 7 days anyway 

144 if i and iweek == 53: 

145 i = 0 

146 iweek = 52 

147 delta = datetime.timedelta(days=7 * i) 

148 

149 # Correct for a weekly being issued in the last week of the 

150 # previous year, as Thursdays don't always line up evenly in 

151 # a week / year split. 

152 dt = datetime.datetime.fromisocalendar(iyear, iweek, dayOfWeek) - delta 

153 # Skip tags that can't be understood. 

154 case _: 

155 continue 

156 times.append(dt) 

157 if len(times) == 0: 

158 raise ValueError("Could not find any tags corresponding to dates") 

159 minTime = min(times) 

160 minTime.replace(tzinfo=datetime.UTC) 

161 return version, minTime.timestamp() 

162 

163 

164@dataclass 

165class SasquatchDispatcher: 

166 """This class mediates the transfer of MetricMeasurementBundles to a 

167 Sasquatch http kafka proxy server. 

168 """ 

169 

170 url: str 

171 """Url of the Sasquatch proxy server""" 

172 

173 token: str 

174 """Authentication token used in communicating with the proxy server""" 

175 

176 namespace: str = "lsst.dm" 

177 """The namespace in Sasquatch in which to write the uploaded metrics""" 

178 

179 def __post_init__(self) -> None: 

180 match ResourcePath(self.url).scheme: 

181 case "http" | "https": 

182 pass 

183 case _: 

184 raise ValueError("Proxy server must be locatable with either http or https") 

185 

186 self._cluster_id: str | None = None 

187 

188 @property 

189 def clusterId(self) -> str: 

190 """ClusterId of the Kafka proxy 

191 

192 Notes 

193 ----- 

194 The cluster Id will be fetched with a network call if it is not 

195 already cached. 

196 """ 

197 if self._cluster_id is None: 

198 self._populateClusterId() 

199 return cast(str, self._cluster_id) 

200 

201 def _populateClusterId(self) -> None: 

202 """Get Sasquatch kafka cluster ID.""" 

203 

204 headers = {"content-type": "application/json"} 

205 

206 try: 

207 with http_client() as session: 

208 r = session.get(f"{self.url}/v3/clusters", headers=headers) 

209 cluster_id = r.json()["data"][0]["cluster_id"] 

210 self._cluster_id = str(cluster_id) 

211 r.raise_for_status() 

212 except requests.RequestException as e: 

213 raise SasquatchDispatchFailure("Could not retrieve the cluster id for the specified url") from e 

214 

215 def _create_topic(self, topic_name: str) -> bool: 

216 """Create a kafka topic in Sasquatch. 

217 

218 Parameters 

219 ---------- 

220 topic_name : `str` 

221 The name of the kafka topic to create 

222 

223 Returns 

224 ------- 

225 status : `bool` 

226 If this does not encounter an error it will return a True success 

227 code, else it will return a False code. 

228 

229 """ 

230 

231 headers = {"content-type": "application/json"} 

232 

233 topic_config = { 

234 "topic_name": f"{self.namespace}.{topic_name}", 

235 "partitions_count": PARTITIONS, 

236 "replication_factor": REPLICATION_FACTOR, 

237 } 

238 

239 try: 

240 with http_client() as session: 

241 r = session.post( 

242 f"{self.url}/v3/clusters/{self.clusterId}/topics", json=topic_config, headers=headers 

243 ) 

244 r.raise_for_status() 

245 log.debug("Created topic %s.%s", self.namespace, topic_name) 

246 return True 

247 except requests.HTTPError as e: 

248 if e.response.status_code == requests.codes.bad_request: 

249 log.debug("Topic %s.%s already exists.", self.namespace, topic_name) 

250 return True 

251 else: 

252 log.error( 

253 "Unknown error occurred creating kafka topic %s %s", 

254 e.response.status_code, 

255 e.response.reason, 

256 ) 

257 return False 

258 except requests.RequestException: 

259 log.exception("Unknown error occurred creating kafka topic") 

260 return False 

261 

262 def _generateAvroSchema(self, metric: str, record: MutableMapping[str, Any]) -> tuple[str, bool]: 

263 """Infer the Avro schema from the record payload. 

264 

265 Parameters 

266 ---------- 

267 metric : `str` 

268 The name of the metric 

269 record : `MutableMapping` 

270 The prepared record for which a schema is to be generated 

271 

272 Returns 

273 ------- 

274 resultSchema : `str` 

275 A json encoded string of the resulting avro schema 

276 errorCode : bool 

277 A boolean indicating if any record fields had to be trimmed because 

278 a suitable schema could not be generated. True if records were 

279 removed, False otherwise. 

280 """ 

281 schema: dict[str, Any] = {"type": "record", "namespace": self.namespace, "name": metric} 

282 

283 # Record if any records needed to be trimmed 

284 resultsTrimmed = False 

285 

286 fields = list() 

287 # If avro schemas cant be generated for values, they should be removed 

288 # from the records. 

289 keysToRemove: list[str] = [] 

290 for key in record: 

291 value = record[key] 

292 avroType: Mapping[str, Any] 

293 if "timestamp" in key: 

294 avroType = {"type": "double"} 

295 else: 

296 avroType = self._python2Avro(value) 

297 if len(avroType) == 0: 

298 continue 

299 if avroType.get("error_in_conversion"): 

300 keysToRemove.append(key) 

301 resultsTrimmed = True 

302 continue 

303 fields.append({"name": key, **avroType}) 

304 

305 # remove any key that failed to have schema generated 

306 for key in keysToRemove: 

307 record.pop(key) 

308 

309 schema["fields"] = fields 

310 

311 return json.dumps(schema), resultsTrimmed 

312 

313 def _python2Avro(self, value: Any) -> Mapping: 

314 """Map python type to avro schema 

315 

316 Parameters 

317 ---------- 

318 value : `Any` 

319 Any python parameter. 

320 

321 Returns 

322 ------- 

323 result : `Mapping` 

324 Return a mapping that represents an entry in an avro schema. 

325 """ 

326 match value: 

327 case float() | np.float32() | None: 

328 return {"type": "float", "default": 0.0} 

329 case str(): 

330 return {"type": "string", "default": ""} 

331 case int(): 

332 return {"type": "int", "default": 0} 

333 case Sequence(): 

334 tmp = {self._python2Avro(item)["type"] for item in value} 

335 if len(tmp) == 0: 

336 return {} 

337 if len(tmp) > 1: 

338 log.error( 

339 "Sequence contains mixed types: %s, must be homogeneous for avro conversion " 

340 "skipping record", 

341 tmp, 

342 ) 

343 return {"error_in_conversion": True} 

344 return {"type": "array", "items": tmp.pop()} 

345 case _: 

346 log.error("Unsupported type %s, skipping record", type(value)) 

347 return {} 

348 

349 def _handleReferencePackage(self, meta: MutableMapping, bundle: MetricMeasurementBundle) -> None: 

350 """Check to see if there is a reference package. 

351 

352 if there is a reference package, determine the datetime associated with 

353 this reference package. Save the package, the version, and the date to 

354 the common metric fields. 

355 

356 Parameters 

357 ---------- 

358 meta : `MutableMapping` 

359 A mapping which corresponds to fields which should be encoded in 

360 all records. 

361 bundle : `MetricMeasurementBundle` 

362 The bundled metrics 

363 """ 

364 package_version, package_timestamp = "", 0.0 

365 if ref_package := getattr(bundle, "reference_package", ""): 

366 ref_package = bundle.reference_package 

367 packages = getEnvironmentPackages(True) 

368 if package_info := packages.get(ref_package): 

369 try: 

370 package_version, package_timestamp = _tag2VersionTime(package_info) 

371 except ValueError: 

372 # Could not extract package timestamp leaving empty 

373 pass 

374 # explicit handle if None was set in the bundle for the package 

375 meta["reference_package"] = ref_package or "" 

376 meta["reference_package_version"] = package_version 

377 meta["reference_package_timestamp"] = package_timestamp 

378 

379 def _handleTimes(self, meta: MutableMapping, bundle: MetricMeasurementBundle, run: str) -> None: 

380 """Add times to the meta fields mapping. 

381 

382 Add all appropriate timestamp fields to the meta field mapping. These 

383 will be added to all records. 

384 

385 This method will also look at the bundle to see if it defines a 

386 preferred time. It so it sets that time as the main time stamp to be 

387 used for this record. 

388 

389 Parameters 

390 ---------- 

391 meta : `MutableMapping` 

392 A mapping which corresponds to fields which should be encoded in 

393 all records. 

394 bundle : `MetricMeasurementBundle` 

395 The bundled metrics 

396 run : `str` 

397 The `~lsst.daf.butler.Butler` collection where the 

398 `MetricMeasurementBundle` is stored. 

399 """ 

400 # Determine the timestamp associated with the run, if someone abused 

401 # the run collection, use the current timestamp 

402 if re.match(r"\d{8}T\d{6}Z", stamp := run.split("/")[-1]): 

403 run_timestamp = datetime.datetime.strptime(stamp, r"%Y%m%dT%H%M%S%z") 

404 else: 

405 run_timestamp = datetime.datetime.now() 

406 meta["run_timestamp"] = run_timestamp.timestamp() 

407 

408 # If the bundle supports supplying timestamps, dispatch on the type 

409 # specified. 

410 if hasattr(bundle, "timestamp_version") and bundle.timestamp_version: 

411 match bundle.timestamp_version: 

412 case "reference_package_timestamp": 

413 if not meta["reference_package_timestamp"]: 

414 log.error("Reference package timestamp is empty, using run_timestamp") 

415 meta["timestamp"] = meta["run_timestamp"] 

416 else: 

417 meta["timestamp"] = meta["reference_package_timestamp"] 

418 case "run_timestamp": 

419 meta["timestamp"] = meta["run_timestamp"] 

420 case "current_timestamp": 

421 timeStamp = datetime.datetime.now() 

422 meta["timestamp"] = timeStamp.timestamp() 

423 case "dataset_timestamp": 

424 log.error("dataset timestamps are not yet supported, run_timestamp will be used") 

425 meta["timestamp"] = meta["run_timestamp"] 

426 case str(value) if "explicit_timestamp" in value: 

427 try: 

428 _, splitTime = value.split(":") 

429 except ValueError as excpt: 

430 raise ValueError( 

431 "Explicit timestamp must be given in the format 'explicit_timestamp:datetime', " 

432 "where datetime is given in the form '%Y%m%dT%H%M%S%z" 

433 ) from excpt 

434 meta["timestamp"] = datetime.datetime.strptime(splitTime, r"%Y%m%dT%H%M%S%z").timestamp() 

435 case _: 

436 log.error( 

437 "Timestamp version %s is not supported, run_timestamp will be used", 

438 bundle.timestamp_version, 

439 ) 

440 meta["timestamp"] = meta["run_timestamp"] 

441 # Default to using the run_timestamp. 

442 else: 

443 meta["timestamp"] = meta["run_timestamp"] 

444 

445 def _handleIdentifier( 

446 self, 

447 meta: MutableMapping, 

448 identifierFields: Mapping[str, Any] | None, 

449 datasetIdentifier: str | None, 

450 bundle: MetricMeasurementBundle, 

451 ) -> None: 

452 """Add an identifier to the meta record mapping. 

453 

454 If the bundle declares a dataset identifier to use add that to the 

455 record, otherwise use 'Generic' as the identifier. If the 

456 datasetIdentifier parameter is specified, that is used instead of 

457 anything specified by the bundle. 

458 

459 This will also add any identifier fields supplied to the meta record 

460 mapping. 

461 

462 Together these values (in addition to the timestamp and topic) should 

463 uniquely identify an upload to the Sasquatch system. 

464 

465 Parameters 

466 ---------- 

467 meta : `MutableMapping` 

468 A mapping which corresponds to fields which should be encoded in 

469 all records. 

470 identifierFields: `Mapping` or `None` 

471 The keys and values in this mapping will be both added as fields 

472 in the record, and used in creating a unique tag for the uploaded 

473 dataset type. I.e. the timestamp, and the tag will be unique, and 

474 each record will belong to one combination of such. 

475 datasetIdentifier : `str` or `None` 

476 A string which will be used in creating unique identifier tags. 

477 bundle : `MetricMeasurementBundle` 

478 The bundle containing metric values to upload. 

479 """ 

480 identifier: str 

481 if datasetIdentifier is not None: 

482 identifier = datasetIdentifier 

483 elif hasattr(bundle, "dataset_identifier") and bundle.dataset_identifier is not None: 

484 identifier = bundle.dataset_identifier 

485 else: 

486 identifier = "Generic" 

487 

488 meta["dataset_tag"] = identifier 

489 

490 if identifierFields is None: 

491 identifierFields = {} 

492 for key in IDENTIFIER_KEYS: 

493 value = identifierFields.get(key, "") 

494 meta[key] = f"{value}" 

495 

496 def _prepareBundle( 

497 self, 

498 bundle: MetricMeasurementBundle, 

499 run: str, 

500 datasetType: str, 

501 timestamp: datetime.datetime | None = None, 

502 id: UUID | None = None, 

503 identifierFields: Mapping | None = None, 

504 datasetIdentifier: str | None = None, 

505 extraFields: Mapping | None = None, 

506 ) -> tuple[Mapping[str, list[Any]], bool]: 

507 """Encode all of the inputs into a format that can be sent to the 

508 kafka proxy server. 

509 

510 Parameters 

511 ---------- 

512 bundle : `MetricMeasurementBundle` 

513 The bundle containing metric values to upload. 

514 run : `str` 

515 The run name to associate with these metric values. If this bundle 

516 is also stored in the butler, this should be the butler run 

517 collection the bundle is stored in the butler. 

518 datasetType : `str` 

519 The dataset type name associated with this 

520 `MetricMeasurementBundle` 

521 timestamp : `datetime.datetime`, optional 

522 The timestamp to be associated with the measurements in the ingress 

523 database. If this value is None, timestamp will be set by the run 

524 time or current time. 

525 id : `UUID`, optional 

526 The UUID of the `MetricMeasurementBundle` within the butler. If 

527 `None`, a new random UUID will be generated so that each record in 

528 Sasquatch will have a unique value. 

529 identifierFields: `Mapping`, optional 

530 The keys and values in this mapping will be both added as fields 

531 in the record, and used in creating a unique tag for the uploaded 

532 dataset type. I.e. the timestamp, and the tag will be unique, and 

533 each record will belong to one combination of such. 

534 datasetIdentifier : `str`, optional 

535 A string which will be used in creating unique identifier tags. 

536 extraFields: `Mapping`, optional 

537 Extra mapping keys and values that will be added as fields to the 

538 dispatched record. 

539 

540 Returns 

541 ------- 

542 result : `Mapping` of `str` to `list` 

543 A mapping of metric name of list of metric measurement records. 

544 status : `bool` 

545 A status boolean indicating if some records had to be skipped due 

546 to a problem parsing the bundle. 

547 """ 

548 if id is None: 

549 id = uuid4() 

550 sid = str(id) 

551 meta: dict[str, Any] = dict() 

552 

553 # Add other associated common fields 

554 meta["id"] = sid 

555 meta["run"] = run 

556 meta["dataset_type"] = datasetType 

557 

558 # Check to see if the bundle declares a reference package 

559 self._handleReferencePackage(meta, bundle) 

560 

561 # Handle the various timestamps that could be associated with a record 

562 self._handleTimes(meta, bundle, run) 

563 

564 # Always use the supplied timestamp if one was passed to use. 

565 if timestamp is not None: 

566 meta["timestamp"] = timestamp.timestamp() 

567 

568 self._handleIdentifier(meta, identifierFields, datasetIdentifier, bundle) 

569 

570 # Add in any other fields that were supplied to the function call. 

571 if extraFields is not None: 

572 meta.update(extraFields) 

573 

574 metricRecords: dict[str, list[Any]] = dict() 

575 

576 # Record if any records needed skipped 

577 resultsTrimmed = False 

578 

579 # Look at each of the metrics in the bundle (name, values) 

580 for metric, measurements in bundle.items(): 

581 # Create a list which will contain the records for each measurement 

582 # associated with metric. 

583 metricRecordList = metricRecords.setdefault(f"{bundle.metricNamePrefix}{metric}", list()) 

584 

585 record: dict[str, Any] = meta.copy() 

586 

587 # loop over each metric measurement within the metric 

588 for measurement in measurements: 

589 # need to extract any tags, package info, etc 

590 note_key = f"{measurement.metric_name.metric}.metric_tags" 

591 record["tags"] = dict(measurement.notes.items()).get(note_key, list()) 

592 

593 # Missing values are replaced by 0 in sasquatch, see RFC-763. 

594 name = "" 

595 value = 0.0 

596 match measurement.json: 

597 case {"metric": name, "value": None}: 

598 pass 

599 case {"metric": name, "value": value}: 

600 if math.isnan(value): 

601 log.error( 

602 "Measurement %s had a value that is a NaN, dispatch will be skipped", 

603 measurement, 

604 ) 

605 resultsTrimmed = True 

606 continue 

607 # JSON will not serialize np.float32, must cast. 

608 if isinstance(value, np.float32): 

609 value = float(value) 

610 pass 

611 case {"value": _}: 

612 log.error("Measurement %s does not contain the key 'metric'", measurement) 

613 resultsTrimmed = True 

614 continue 

615 case {"metric": _}: 

616 log.error("Measurement %s does not contain the key 'value'", measurement) 

617 resultsTrimmed = True 

618 continue 

619 record[name] = value 

620 

621 metricRecordList.append({"value": record}) 

622 return metricRecords, resultsTrimmed 

623 

624 def dispatch( 

625 self, 

626 bundle: MetricMeasurementBundle, 

627 run: str, 

628 datasetType: str, 

629 timestamp: datetime.datetime | None = None, 

630 id: UUID | None = None, 

631 datasetIdentifier: str | None = None, 

632 identifierFields: Mapping | None = None, 

633 extraFields: Mapping | None = None, 

634 ) -> None: 

635 """Dispatch a `MetricMeasurementBundle` to Sasquatch. 

636 

637 Parameters 

638 ---------- 

639 bundle : `MetricMeasurementBundle` 

640 The bundle containing metric values to upload. 

641 run : `str` 

642 The run name to associate with these metric values. If this bundle 

643 is also stored in the butler, this should be the butler run 

644 collection the bundle is stored in the butler. This will be used 

645 in generating uniqueness constraints in Sasquatch. 

646 datasetType : `str` 

647 The dataset type name associated with this 

648 `MetricMeasurementBundle`. 

649 timestamp : `datetime.datetime`, optional 

650 The timestamp to be associated with the measurements in the ingress 

651 database. If this value is None, timestamp will be set by the run 

652 time or current time. 

653 id : `UUID`, optional 

654 The UUID of the `MetricMeasurementBundle` within the Butler. If 

655 `None`, a new random UUID will be generated so that each record in 

656 Sasquatch will have a unique value. 

657 datasetIdentifier : `str`, optional 

658 A string which will be used in creating unique identifier tags. If 

659 `None`, a default value will be inserted. 

660 identifierFields: `Mapping`, optional 

661 The keys and values in this mapping will be both added as fields 

662 in the record, and used in creating a unique tag for the uploaded 

663 dataset type. I.e. the timestamp, and the tag will be unique, and 

664 each record will belong to one combination of such. Examples of 

665 entries would be things like visit or tract. 

666 extraFields: `Mapping`, optional 

667 Extra mapping keys and values that will be added as fields to the 

668 dispatched record. 

669 

670 Raises 

671 ------ 

672 SasquatchDispatchPartialFailure, SasquatchDispatchFailure 

673 Raised if there were any errors in dispatching a bundle. 

674 """ 

675 if id is None: 

676 id = uuid4() 

677 

678 # Prepare the bundle by transforming it to a list of metric records 

679 metricRecords, recordsTrimmed = self._prepareBundle( 

680 bundle=bundle, 

681 run=run, 

682 datasetType=datasetType, 

683 timestamp=timestamp, 

684 id=id, 

685 datasetIdentifier=datasetIdentifier, 

686 identifierFields=identifierFields, 

687 extraFields=extraFields, 

688 ) 

689 

690 headers = {"content-type": "application/vnd.kafka.avro.v2+json"} 

691 data: dict[str, Any] = dict() 

692 partialUpload = False 

693 uploadFailed = [] 

694 

695 with http_client() as session: 

696 for metric, record in metricRecords.items(): 

697 # create the kafka topic if it does not already exist 

698 if not self._create_topic(metric): 

699 log.error("Topic not created, skipping dispatch of %s", metric) 

700 continue 

701 recordValue = record[0]["value"] 

702 # Generate schemas for each record 

703 data["value_schema"], schemaTrimmed = self._generateAvroSchema(metric, recordValue) 

704 data["records"] = record 

705 

706 if schemaTrimmed: 

707 partialUpload = True 

708 

709 try: 

710 r = session.post( 

711 f"{self.url}/topics/{self.namespace}.{metric}", json=data, headers=headers 

712 ) 

713 r.raise_for_status() 

714 log.debug("Succesfully sent data for metric %s", metric) 

715 uploadFailed.append(False) 

716 except requests.HTTPError as e: 

717 log.error( 

718 "There was a problem submitting the metric %s: %s, %s", 

719 metric, 

720 e.response.status_code, 

721 e.response.reason, 

722 ) 

723 uploadFailed.append(True) 

724 partialUpload = True 

725 except requests.RequestException as e: 

726 # Don't log full stack trace because there may be lots 

727 # of these. 

728 log.error("There was a problem submitting the metric %s: %s", metric, e) 

729 uploadFailed.append(True) 

730 partialUpload = True 

731 

732 # There may be no metrics to try to upload, and thus the uploadFailed 

733 # list may be empty, check before issuing failure 

734 if len(uploadFailed) > 0 and all(uploadFailed): 

735 raise SasquatchDispatchFailure("All records were unable to be uploaded.") 

736 

737 if partialUpload or recordsTrimmed: 

738 raise SasquatchDispatchPartialFailure("One or more records may not have been uploaded entirely") 

739 

740 def dispatchRef( 

741 self, 

742 bundle: MetricMeasurementBundle, 

743 ref: DatasetRef, 

744 timestamp: datetime.datetime | None = None, 

745 extraFields: Mapping | None = None, 

746 datasetIdentifier: str | None = None, 

747 ) -> None: 

748 """Dispatch a `MetricMeasurementBundle` to Sasquatch with a known 

749 `DatasetRef`. 

750 

751 Parameters 

752 ---------- 

753 bundle : `MetricMeasurementBundle` 

754 The bundle containing metric values to upload. 

755 ref : `DatasetRef` 

756 The `Butler` dataset ref corresponding to the input 

757 `MetricMeasurementBundle`. 

758 timestamp : `datetime.datetime`, optional 

759 The timestamp to be associated with the measurements in the ingress 

760 database. If this value is None, timestamp will be set by the run 

761 time or current time. 

762 extraFields: `Mapping`, optional 

763 Extra mapping keys and values that will be added as fields to the 

764 dispatched record if not None. 

765 datasetIdentifier : `str`, optional 

766 A string which will be used in creating unique identifier tags. If 

767 None, a default value will be inserted. 

768 

769 Raises 

770 ------ 

771 SasquatchDispatchPartialFailure, SasquatchDispatchFailure 

772 Raised if there were any errors in dispatching a bundle. 

773 """ 

774 # Parse the relevant info out of the dataset ref. 

775 serializedRef = ref.to_simple() 

776 id = serializedRef.id 

777 if serializedRef.run is None: 

778 run = "<unknown>" 

779 else: 

780 run = serializedRef.run 

781 dstype = serializedRef.datasetType 

782 datasetType = dstype.name if dstype is not None else "" 

783 dataRefMapping = serializedRef.dataId.dataId if serializedRef.dataId else None 

784 

785 self.dispatch( 

786 bundle, 

787 run=run, 

788 timestamp=timestamp, 

789 datasetType=datasetType, 

790 id=id, 

791 identifierFields=dataRefMapping, 

792 extraFields=extraFields, 

793 datasetIdentifier=datasetIdentifier, 

794 )