Coverage for python / lsst / dax / apdb / tests / _apdb.py: 15%

1# This file is part of dax_apdb. 

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__ = ["ApdbSchemaUpdateTest", "ApdbTest", "update_schema_yaml"] 

25 

26import contextlib 

27import logging.config 

28import os 

29import tempfile 

30from abc import ABC, abstractmethod 

31from collections.abc import Iterator 

32from tempfile import TemporaryDirectory 

33from typing import TYPE_CHECKING, Any 

34 

35import astropy.time 

36import felis.datamodel 

37import pandas 

38import yaml 

39 

40from lsst.sphgeom import Angle, Circle, LonLat, Region, UnitVector3d 

41 

42from .. import ( 

43 Apdb, 

44 ApdbConfig, 

45 ApdbReassignDiaSourceRecord, 

46 ApdbReplica, 

47 ApdbTableData, 

48 ApdbTables, 

49 ApdbUpdateRecord, 

50 ApdbWithdrawDiaSourceRecord, 

51 IncompatibleVersionError, 

52 ReplicaChunk, 

53 VersionTuple, 

54) 

55from .data_factory import ( 

56 makeForcedSourceCatalog, 

57 makeObjectCatalog, 

58 makeSourceCatalog, 

59 makeTimestampNow, 

60) 

61from .utils import TestCaseMixin 

62 

63if TYPE_CHECKING: 

64 from ..pixelization import Pixelization 

65 

66 

67# Optionally configure logging from a config file. 

68if log_config := os.environ.get("DAX_APDB_TEST_LOG_CONFIG"): 68 ↛ 69line 68 didn't jump to line 69 because the condition on line 68 was never true

69 logging.config.fileConfig(log_config) 

70 

71 

72def _make_region(xyz: tuple[float, float, float] = (1.0, 1.0, -1.0)) -> Region: 

73 """Make a region to use in tests""" 

74 pointing_v = UnitVector3d(*xyz) 

75 fov = 0.0013 # radians 

76 region = Circle(pointing_v, Angle(fov / 2)) 

77 return region 

78 

79 

80@contextlib.contextmanager 

81def update_schema_yaml( 

82 schema_file: str, 

83 drop_metadata: bool = False, 

84 version: str | None = None, 

85) -> Iterator[str]: 

86 """Update schema definition and return name of the new schema file. 

87 

88 Parameters 

89 ---------- 

90 schema_file : `str` 

91 Path for the existing YAML file with APDB schema. 

92 drop_metadata : `bool` 

93 If `True` then remove metadata table from the list of tables. 

94 version : `str` or `None` 

95 If non-empty string then set schema version to this string, if empty 

96 string then remove schema version from config, if `None` - don't change 

97 the version in config. 

98 

99 Yields 

100 ------ 

101 Path for the updated configuration file. 

102 """ 

103 with open(schema_file) as yaml_stream: 

104 schemas_list = list(yaml.load_all(yaml_stream, Loader=yaml.SafeLoader)) 

105 # Edit YAML contents. 

106 for schema in schemas_list: 

107 # Optionally drop metadata table. 

108 if drop_metadata: 

109 schema["tables"] = [table for table in schema["tables"] if table["name"] != "metadata"] 

110 if version is not None: 

111 if version == "": 

112 del schema["version"] 

113 else: 

114 schema["version"] = version 

115 

116 with TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir: 

117 output_path = os.path.join(tmpdir, "schema.yaml") 

118 with open(output_path, "w") as yaml_stream: 

119 yaml.dump_all(schemas_list, stream=yaml_stream) 

120 yield output_path 

121 

122 

123class ApdbTest(TestCaseMixin, ABC): 

124 """Base class for Apdb tests that can be specialized for concrete 

125 implementation. 

126 

127 This can only be used as a mixin class for a unittest.TestCase and it 

128 calls various assert methods. 

129 """ 

130 

131 visit_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai") 

132 

133 fsrc_requires_id_list = False 

134 """Should be set to True if getDiaForcedSources requires object IDs""" 

135 

136 enable_replica: bool = False 

137 """Set to true when support for replication is configured""" 

138 

139 use_mjd: bool = True 

140 """If True then timestamp columns are MJD TAI.""" 

141 

142 extra_chunk_columns = 1 

143 """Number of additional columns in chunk tables.""" 

144 

145 meta_row_count = 3 

146 """Initial row count in metadata table.""" 

147 

148 # number of columns as defined in tests/config/schema.yaml 

149 table_column_count = { 

150 ApdbTables.DiaObject: 7, 

151 ApdbTables.DiaObjectLast: 5, 

152 ApdbTables.DiaSource: 12, 

153 ApdbTables.DiaForcedSource: 8, 

154 ApdbTables.SSObject: 3, 

155 } 

156 

157 @abstractmethod 

158 def make_instance(self, **kwargs: Any) -> ApdbConfig: 

159 """Make database instance and return configuration for it.""" 

160 raise NotImplementedError() 

161 

162 @abstractmethod 

163 def getDiaObjects_table(self) -> ApdbTables: 

164 """Return type of table returned from getDiaObjects method.""" 

165 raise NotImplementedError() 

166 

167 @abstractmethod 

168 def pixelization(self, config: ApdbConfig) -> Pixelization: 

169 """Return pixelization used by implementation.""" 

170 raise NotImplementedError() 

171 

172 def assert_catalog(self, catalog: Any, rows: int, table: ApdbTables) -> None: 

173 """Validate catalog type and size 

174 

175 Parameters 

176 ---------- 

177 catalog : `object` 

178 Expected type of this is ``pandas.DataFrame``. 

179 rows : `int` 

180 Expected number of rows in a catalog. 

181 table : `ApdbTables` 

182 APDB table type. 

183 """ 

184 self.assertIsInstance(catalog, pandas.DataFrame) 

185 self.assertEqual(catalog.shape[0], rows) 

186 self.assertEqual(catalog.shape[1], self.table_column_count[table]) 

187 

188 def assert_table_data(self, catalog: Any, rows: int, table: ApdbTables) -> None: 

189 """Validate catalog type and size 

190 

191 Parameters 

192 ---------- 

193 catalog : `object` 

194 Expected type of this is `ApdbTableData`. 

195 rows : `int` 

196 Expected number of rows in a catalog. 

197 table : `ApdbTables` 

198 APDB table type. 

199 extra_columns : `int` 

200 Count of additional columns expected in ``catalog``. 

201 """ 

202 self.assertIsInstance(catalog, ApdbTableData) 

203 n_rows = sum(1 for row in catalog.rows()) 

204 self.assertEqual(n_rows, rows) 

205 # One extra column for replica chunk id 

206 self.assertEqual( 

207 len(catalog.column_names()), self.table_column_count[table] + self.extra_chunk_columns 

208 ) 

209 

210 def assert_column_types(self, catalog: Any, types: dict[str, felis.datamodel.DataType]) -> None: 

211 column_defs = dict(catalog.column_defs()) 

212 for column, datatype in types.items(): 

213 self.assertEqual(column_defs[column], datatype) 

214 

215 def make_region(self, xyz: tuple[float, float, float] = (1.0, 1.0, -1.0)) -> Region: 

216 """Make a region to use in tests""" 

217 return _make_region(xyz) 

218 

219 def test_makeSchema(self) -> None: 

220 """Test for making APDB schema.""" 

221 config = self.make_instance() 

222 apdb = Apdb.from_config(config) 

223 

224 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject)) 

225 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast)) 

226 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource)) 

227 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource)) 

228 self.assertIsNotNone(apdb.tableDef(ApdbTables.metadata)) 

229 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSObject)) 

230 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSSource)) 

231 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject_To_Object_Match)) 

232 

233 # Test from_uri factory method with the same config. 

234 with tempfile.NamedTemporaryFile() as tmpfile: 

235 config.save(tmpfile.name) 

236 apdb = Apdb.from_uri(tmpfile.name) 

237 

238 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject)) 

239 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast)) 

240 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource)) 

241 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource)) 

242 self.assertIsNotNone(apdb.tableDef(ApdbTables.metadata)) 

243 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSObject)) 

244 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSSource)) 

245 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject_To_Object_Match)) 

246 

247 def test_empty_gets(self) -> None: 

248 """Test for getting data from empty database. 

249 

250 All get() methods should return empty results, only useful for 

251 checking that code is not broken. 

252 """ 

253 # use non-zero months for Forced/Source fetching 

254 config = self.make_instance() 

255 apdb = Apdb.from_config(config) 

256 

257 region = self.make_region() 

258 visit_time = self.visit_time 

259 

260 res: pandas.DataFrame | None 

261 

262 # get objects by region 

263 res = apdb.getDiaObjects(region) 

264 self.assert_catalog(res, 0, self.getDiaObjects_table()) 

265 

266 # get sources by region 

267 res = apdb.getDiaSources(region, None, visit_time) 

268 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

269 

270 res = apdb.getDiaSources(region, [], visit_time) 

271 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

272 

273 # get sources by object ID, non-empty object list 

274 res = apdb.getDiaSources(region, [1, 2, 3], visit_time) 

275 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

276 

277 # get forced sources by object ID, empty object list 

278 res = apdb.getDiaForcedSources(region, [], visit_time) 

279 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

280 

281 # get sources by object ID, non-empty object list 

282 res = apdb.getDiaForcedSources(region, [1, 2, 3], visit_time) 

283 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

284 

285 # data_factory's ccdVisitId generation corresponds to (1, 1) 

286 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time) 

287 self.assertFalse(res) 

288 

289 # get sources by region 

290 if self.fsrc_requires_id_list: 

291 with self.assertRaises(NotImplementedError): 

292 apdb.getDiaForcedSources(region, None, visit_time) 

293 else: 

294 res = apdb.getDiaForcedSources(region, None, visit_time) 

295 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

296 

297 def test_empty_gets_0months(self) -> None: 

298 """Test for getting data from empty database. 

299 

300 All get() methods should return empty DataFrame or None. 

301 """ 

302 # set read_sources_months to 0 so that Forced/Sources are None 

303 config = self.make_instance(read_sources_months=0, read_forced_sources_months=0) 

304 apdb = Apdb.from_config(config) 

305 

306 region = self.make_region() 

307 visit_time = self.visit_time 

308 

309 res: pandas.DataFrame | None 

310 

311 # get objects by region 

312 res = apdb.getDiaObjects(region) 

313 self.assert_catalog(res, 0, self.getDiaObjects_table()) 

314 

315 # get sources by region 

316 res = apdb.getDiaSources(region, None, visit_time) 

317 self.assertIs(res, None) 

318 

319 # get sources by object ID, empty object list 

320 res = apdb.getDiaSources(region, [], visit_time) 

321 self.assertIs(res, None) 

322 

323 # get forced sources by object ID, empty object list 

324 res = apdb.getDiaForcedSources(region, [], visit_time) 

325 self.assertIs(res, None) 

326 

327 # Database is empty, no images exist. 

328 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time) 

329 self.assertFalse(res) 

330 

331 def test_storeObjects(self) -> None: 

332 """Store and retrieve DiaObjects.""" 

333 # don't care about sources. 

334 config = self.make_instance() 

335 apdb = Apdb.from_config(config) 

336 

337 region = self.make_region() 

338 visit_time = self.visit_time 

339 

340 # make catalog with Objects 

341 catalog = makeObjectCatalog(region, 100, visit_time) 

342 

343 # store catalog 

344 apdb.store(visit_time, catalog) 

345 

346 # read it back and check sizes 

347 res = apdb.getDiaObjects(region) 

348 self.assert_catalog(res, len(catalog), self.getDiaObjects_table()) 

349 

350 # TODO: test apdb.contains with generic implementation from DM-41671 

351 

352 def test_storeObjects_empty(self) -> None: 

353 """Test calling storeObject when there are no objects: see DM-43270.""" 

354 config = self.make_instance() 

355 apdb = Apdb.from_config(config) 

356 region = self.make_region() 

357 visit_time = self.visit_time 

358 # make catalog with no Objects 

359 catalog = makeObjectCatalog(region, 0, visit_time) 

360 

361 with self.assertLogs("lsst.dax.apdb", level="DEBUG") as cm: 

362 apdb.store(visit_time, catalog) 

363 self.assertIn("No objects", "\n".join(cm.output)) 

364 

365 def test_storeMovingObject(self) -> None: 

366 """Store and retrieve DiaObject which changes its position.""" 

367 # don't care about sources. 

368 config = self.make_instance() 

369 apdb = Apdb.from_config(config) 

370 pixelization = self.pixelization(config) 

371 

372 lon_deg, lat_deg = 0.0, 0.0 

373 lonlat1 = LonLat.fromDegrees(lon_deg - 1.0, lat_deg) 

374 lonlat2 = LonLat.fromDegrees(lon_deg + 1.0, lat_deg) 

375 uv1 = UnitVector3d(lonlat1) 

376 uv2 = UnitVector3d(lonlat2) 

377 

378 # Check that they fall into different pixels. 

379 self.assertNotEqual(pixelization.pixel(uv1), pixelization.pixel(uv2)) 

380 

381 # Store one object at two different positions. 

382 visit_time1 = self.visit_time 

383 catalog1 = makeObjectCatalog(lonlat1, 1, visit_time1) 

384 apdb.store(visit_time1, catalog1) 

385 

386 visit_time2 = visit_time1 + astropy.time.TimeDelta(120.0, format="sec") 

387 catalog1 = makeObjectCatalog(lonlat2, 1, visit_time2) 

388 apdb.store(visit_time2, catalog1) 

389 

390 # Make region covering both points. 

391 region = Circle(UnitVector3d(LonLat.fromDegrees(lon_deg, lat_deg)), Angle.fromDegrees(1.1)) 

392 self.assertTrue(region.contains(uv1)) 

393 self.assertTrue(region.contains(uv2)) 

394 

395 # Read it back, must return the latest one. 

396 res = apdb.getDiaObjects(region) 

397 self.assert_catalog(res, 1, self.getDiaObjects_table()) 

398 

399 def test_storeSources(self) -> None: 

400 """Store and retrieve DiaSources.""" 

401 config = self.make_instance() 

402 apdb = Apdb.from_config(config) 

403 

404 region = self.make_region() 

405 visit_time = self.visit_time 

406 

407 # have to store Objects first 

408 objects = makeObjectCatalog(region, 100, visit_time) 

409 oids = list(objects["diaObjectId"]) 

410 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd) 

411 

412 # save the objects and sources 

413 apdb.store(visit_time, objects, sources) 

414 

415 # read it back, no ID filtering 

416 res = apdb.getDiaSources(region, None, visit_time) 

417 self.assert_catalog(res, len(sources), ApdbTables.DiaSource) 

418 

419 # read it back and filter by ID 

420 res = apdb.getDiaSources(region, oids, visit_time) 

421 self.assert_catalog(res, len(sources), ApdbTables.DiaSource) 

422 

423 # read it back to get schema 

424 res = apdb.getDiaSources(region, [], visit_time) 

425 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

426 

427 # test if a visit is present 

428 # data_factory's ccdVisitId generation corresponds to (1, 1) 

429 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time) 

430 self.assertTrue(res) 

431 # non-existent image 

432 res = apdb.containsVisitDetector(visit=2, detector=42, region=region, visit_time=visit_time) 

433 self.assertFalse(res) 

434 

435 def test_storeForcedSources(self) -> None: 

436 """Store and retrieve DiaForcedSources.""" 

437 config = self.make_instance() 

438 apdb = Apdb.from_config(config) 

439 

440 region = self.make_region() 

441 visit_time = self.visit_time 

442 

443 # have to store Objects first 

444 objects = makeObjectCatalog(region, 100, visit_time) 

445 oids = list(objects["diaObjectId"]) 

446 catalog = makeForcedSourceCatalog(objects, visit_time, use_mjd=self.use_mjd) 

447 

448 apdb.store(visit_time, objects, forced_sources=catalog) 

449 

450 # read it back and check sizes 

451 res = apdb.getDiaForcedSources(region, oids, visit_time) 

452 self.assert_catalog(res, len(catalog), ApdbTables.DiaForcedSource) 

453 

454 # read it back to get schema 

455 res = apdb.getDiaForcedSources(region, [], visit_time) 

456 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource) 

457 

458 # data_factory's ccdVisitId generation corresponds to (1, 1) 

459 res = apdb.containsVisitDetector(visit=1, detector=1, region=region, visit_time=visit_time) 

460 self.assertTrue(res) 

461 # non-existent image 

462 res = apdb.containsVisitDetector(visit=2, detector=42, region=region, visit_time=visit_time) 

463 self.assertFalse(res) 

464 

465 def test_timestamps(self) -> None: 

466 """Check that timestamp return type is as expected.""" 

467 config = self.make_instance() 

468 apdb = Apdb.from_config(config) 

469 

470 region = self.make_region() 

471 visit_time = self.visit_time 

472 

473 # Cassandra has a millisecond precision, so subtract 1ms to allow for 

474 # truncated returned values. 

475 time_before = makeTimestampNow(self.use_mjd, -1) 

476 objects = makeObjectCatalog(region, 100, visit_time) 

477 oids = list(objects["diaObjectId"]) 

478 catalog = makeForcedSourceCatalog(objects, visit_time, use_mjd=self.use_mjd) 

479 time_after = makeTimestampNow(self.use_mjd) 

480 

481 apdb.store(visit_time, objects, forced_sources=catalog) 

482 

483 # read it back and check sizes 

484 res = apdb.getDiaForcedSources(region, oids, visit_time) 

485 assert res is not None 

486 self.assert_catalog(res, len(catalog), ApdbTables.DiaForcedSource) 

487 

488 time_processed_column = "timeProcessedMjdTai" if self.use_mjd else "time_processed" 

489 self.assertIn(time_processed_column, res.dtypes) 

490 dtype = res.dtypes[time_processed_column] 

491 timestamp_type_name = "float64" if self.use_mjd else "datetime64[ns]" 

492 self.assertEqual(dtype.name, timestamp_type_name) 

493 # Verify that returned time is sensible. 

494 self.assertTrue(all(time_before <= dt <= time_after for dt in res[time_processed_column])) 

495 

496 def test_getChunks(self) -> None: 

497 """Store and retrieve replica chunks.""" 

498 # don't care about sources. 

499 config = self.make_instance() 

500 apdb = Apdb.from_config(config) 

501 apdb_replica = ApdbReplica.from_config(config) 

502 visit_time = self.visit_time 

503 

504 region1 = self.make_region((1.0, 1.0, -1.0)) 

505 region2 = self.make_region((-1.0, -1.0, -1.0)) 

506 nobj = 100 

507 objects1 = makeObjectCatalog(region1, nobj, visit_time) 

508 objects2 = makeObjectCatalog(region2, nobj, visit_time, start_id=nobj * 2) 

509 

510 # With the default 10 minutes replica chunk window we should have 4 

511 # records. 

512 visits = [ 

513 (astropy.time.Time("2021-01-01T00:01:00", format="isot", scale="tai"), objects1), 

514 (astropy.time.Time("2021-01-01T00:02:00", format="isot", scale="tai"), objects2), 

515 (astropy.time.Time("2021-01-01T00:11:00", format="isot", scale="tai"), objects1), 

516 (astropy.time.Time("2021-01-01T00:12:00", format="isot", scale="tai"), objects2), 

517 (astropy.time.Time("2021-01-01T00:45:00", format="isot", scale="tai"), objects1), 

518 (astropy.time.Time("2021-01-01T00:46:00", format="isot", scale="tai"), objects2), 

519 (astropy.time.Time("2021-03-01T00:01:00", format="isot", scale="tai"), objects1), 

520 (astropy.time.Time("2021-03-01T00:02:00", format="isot", scale="tai"), objects2), 

521 ] 

522 

523 start_id = 0 

524 for visit_time, objects in visits: 

525 sources = makeSourceCatalog(objects, visit_time, start_id=start_id, use_mjd=self.use_mjd) 

526 fsources = makeForcedSourceCatalog(objects, visit_time, visit=start_id, use_mjd=self.use_mjd) 

527 apdb.store(visit_time, objects, sources, fsources) 

528 start_id += nobj 

529 

530 replica_chunks = apdb_replica.getReplicaChunks() 

531 if not self.enable_replica: 

532 self.assertIsNone(replica_chunks) 

533 

534 with self.assertRaisesRegex(ValueError, "APDB is not configured for replication"): 

535 apdb_replica.getTableDataChunks(ApdbTables.DiaObject, []) 

536 

537 else: 

538 assert replica_chunks is not None 

539 self.assertEqual(len(replica_chunks), 4) 

540 

541 with self.assertRaisesRegex(ValueError, "does not support replica chunks"): 

542 apdb_replica.getTableDataChunks(ApdbTables.SSObject, []) 

543 

544 def _check_chunks(replica_chunks: list[ReplicaChunk], n_records: int | None = None) -> None: 

545 if n_records is None: 

546 n_records = len(replica_chunks) * nobj 

547 res = apdb_replica.getTableDataChunks( 

548 ApdbTables.DiaObject, (chunk.id for chunk in replica_chunks) 

549 ) 

550 self.assert_table_data(res, n_records, ApdbTables.DiaObject) 

551 validityStartColumn = "validityStartMjdTai" if self.use_mjd else "validityStart" 

552 validityStartType = ( 

553 felis.datamodel.DataType.double if self.use_mjd else felis.datamodel.DataType.timestamp 

554 ) 

555 self.assert_column_types( 

556 res, 

557 { 

558 "apdb_replica_chunk": felis.datamodel.DataType.long, 

559 "diaObjectId": felis.datamodel.DataType.long, 

560 validityStartColumn: validityStartType, 

561 "ra": felis.datamodel.DataType.double, 

562 "dec": felis.datamodel.DataType.double, 

563 "parallax": felis.datamodel.DataType.float, 

564 "nDiaSources": felis.datamodel.DataType.int, 

565 }, 

566 ) 

567 

568 res = apdb_replica.getTableDataChunks( 

569 ApdbTables.DiaSource, (chunk.id for chunk in replica_chunks) 

570 ) 

571 self.assert_table_data(res, n_records, ApdbTables.DiaSource) 

572 self.assert_column_types( 

573 res, 

574 { 

575 "apdb_replica_chunk": felis.datamodel.DataType.long, 

576 "diaSourceId": felis.datamodel.DataType.long, 

577 "visit": felis.datamodel.DataType.long, 

578 "detector": felis.datamodel.DataType.short, 

579 }, 

580 ) 

581 

582 res = apdb_replica.getTableDataChunks( 

583 ApdbTables.DiaForcedSource, (chunk.id for chunk in replica_chunks) 

584 ) 

585 self.assert_table_data(res, n_records, ApdbTables.DiaForcedSource) 

586 self.assert_column_types( 

587 res, 

588 { 

589 "apdb_replica_chunk": felis.datamodel.DataType.long, 

590 "diaObjectId": felis.datamodel.DataType.long, 

591 "visit": felis.datamodel.DataType.long, 

592 "detector": felis.datamodel.DataType.short, 

593 }, 

594 ) 

595 

596 # read it back and check sizes 

597 _check_chunks(replica_chunks, 800) 

598 _check_chunks(replica_chunks[1:], 600) 

599 _check_chunks(replica_chunks[1:-1], 400) 

600 _check_chunks(replica_chunks[2:3], 200) 

601 _check_chunks([]) 

602 

603 # try to remove some of those 

604 deleted_chunks = replica_chunks[:1] 

605 apdb_replica.deleteReplicaChunks(chunk.id for chunk in deleted_chunks) 

606 

607 # All queries on deleted ids should return empty set. 

608 _check_chunks(deleted_chunks, 0) 

609 

610 replica_chunks = apdb_replica.getReplicaChunks() 

611 assert replica_chunks is not None 

612 self.assertEqual(len(replica_chunks), 3) 

613 

614 _check_chunks(replica_chunks, 600) 

615 

616 def test_reassignObjects(self) -> None: 

617 """Reassign DiaObjects.""" 

618 # don't care about sources. 

619 config = self.make_instance() 

620 apdb = Apdb.from_config(config) 

621 

622 region = self.make_region() 

623 visit_time = self.visit_time 

624 objects = makeObjectCatalog(region, 100, visit_time) 

625 oids = list(objects["diaObjectId"]) 

626 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd) 

627 apdb.store(visit_time, objects, sources) 

628 

629 # read it back and filter by ID 

630 res = apdb.getDiaSources(region, oids, visit_time) 

631 self.assert_catalog(res, len(sources), ApdbTables.DiaSource) 

632 

633 apdb.reassignDiaSources({1: 1, 2: 2, 5: 5}) 

634 res = apdb.getDiaSources(region, oids, visit_time) 

635 self.assert_catalog(res, len(sources) - 3, ApdbTables.DiaSource) 

636 

637 with self.assertRaisesRegex(ValueError, r"do not exist.*\D1000"): 

638 apdb.reassignDiaSources( 

639 { 

640 1000: 1, 

641 7: 3, 

642 } 

643 ) 

644 self.assert_catalog(res, len(sources) - 3, ApdbTables.DiaSource) 

645 

646 def test_storeUpdateRecord(self) -> None: 

647 """Test _storeUpdateRecord() method.""" 

648 config = self.make_instance() 

649 apdb = Apdb.from_config(config) 

650 

651 # Times are totally arbitrary. 

652 update_time_ns1 = 2_000_000_000_000_000_000 

653 update_time_ns2 = 2_000_000_001_000_000_000 

654 records = [ 

655 ApdbReassignDiaSourceRecord( 

656 update_time_ns=update_time_ns1, 

657 update_order=0, 

658 diaSourceId=1, 

659 diaObjectId=321, 

660 ssObjectId=1, 

661 ssObjectReassocTimeMjdTai=60000.0, 

662 ra=45.0, 

663 dec=-45.0, 

664 ), 

665 ApdbWithdrawDiaSourceRecord( 

666 update_time_ns=update_time_ns1, 

667 update_order=1, 

668 diaSourceId=123456, 

669 diaObjectId=321, 

670 timeWithdrawnMjdTai=61000.0, 

671 ra=45.0, 

672 dec=-45.0, 

673 ), 

674 ApdbReassignDiaSourceRecord( 

675 update_time_ns=update_time_ns1, 

676 update_order=3, 

677 diaSourceId=2, 

678 diaObjectId=3, 

679 ssObjectId=3, 

680 ssObjectReassocTimeMjdTai=60000.0, 

681 ra=45.0, 

682 dec=-45.0, 

683 ), 

684 ApdbWithdrawDiaSourceRecord( 

685 update_time_ns=update_time_ns2, 

686 update_order=0, 

687 diaSourceId=123456, 

688 diaObjectId=321, 

689 timeWithdrawnMjdTai=61000.0, 

690 ra=45.0, 

691 dec=-45.0, 

692 ), 

693 ] 

694 

695 update_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai") 

696 chunk = ReplicaChunk.make_replica_chunk(update_time, 600) 

697 

698 if not self.enable_replica: 

699 with self.assertRaises(TypeError): 

700 self.store_update_records(apdb, records, chunk) 

701 else: 

702 self.store_update_records(apdb, records, chunk) 

703 

704 apdb_replica = ApdbReplica.from_config(config) 

705 records_returned = apdb_replica.getUpdateRecordChunks([chunk.id]) 

706 

707 # Input records are ordered, output will be ordered too. 

708 self.assertEqual(records_returned, records) 

709 

710 @abstractmethod 

711 def store_update_records(self, apdb: Apdb, records: list[ApdbUpdateRecord], chunk: ReplicaChunk) -> None: 

712 """Store update records in database, must be overriden in subclass.""" 

713 raise NotImplementedError() 

714 

715 def test_midpointMjdTai_src(self) -> None: 

716 """Test for time filtering of DiaSources.""" 

717 config = self.make_instance() 

718 apdb = Apdb.from_config(config) 

719 

720 region = self.make_region() 

721 # 2021-01-01 plus 360 days is 2021-12-27 

722 src_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai") 

723 src_time2 = astropy.time.Time("2021-01-01T00:00:02", format="isot", scale="tai") 

724 visit_time0 = astropy.time.Time("2021-12-26T23:59:59", format="isot", scale="tai") 

725 visit_time1 = astropy.time.Time("2021-12-27T00:00:01", format="isot", scale="tai") 

726 visit_time2 = astropy.time.Time("2021-12-27T00:00:03", format="isot", scale="tai") 

727 one_sec = astropy.time.TimeDelta(1.0, format="sec") 

728 

729 objects = makeObjectCatalog(region, 100, visit_time0) 

730 oids = list(objects["diaObjectId"]) 

731 sources = makeSourceCatalog(objects, src_time1, 0, use_mjd=self.use_mjd) 

732 apdb.store(src_time1, objects, sources) 

733 

734 sources = makeSourceCatalog(objects, src_time2, 100, use_mjd=self.use_mjd) 

735 apdb.store(src_time2, objects, sources) 

736 

737 # reading at time of last save should read all 

738 res = apdb.getDiaSources(region, oids, src_time2) 

739 self.assert_catalog(res, 200, ApdbTables.DiaSource) 

740 

741 # one second before 12 months 

742 res = apdb.getDiaSources(region, oids, visit_time0) 

743 self.assert_catalog(res, 200, ApdbTables.DiaSource) 

744 

745 # reading at later time of last save should only read a subset 

746 res = apdb.getDiaSources(region, oids, visit_time1) 

747 self.assert_catalog(res, 100, ApdbTables.DiaSource) 

748 

749 # reading at later time of last save should only read a subset 

750 res = apdb.getDiaSources(region, oids, visit_time2) 

751 self.assert_catalog(res, 0, ApdbTables.DiaSource) 

752 

753 # Use explicit start time argument instead of 12 month window, visit 

754 # time does not matter in this case, set it to before all data. 

755 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time1 - one_sec) 

756 self.assert_catalog(res, 200, ApdbTables.DiaSource) 

757 

758 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time2 - one_sec)