Coverage for tests/test_apdbSql.py: 94%

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 

22"""Unit test for Apdb class.""" 

23 

24import gc 

25import os 

26import shutil 

27import tempfile 

28import unittest 

29from contextlib import closing 

30from typing import Any 

31from unittest.mock import patch 

32 

33import felis.datamodel 

34import sqlalchemy 

35 

36import lsst.utils.tests 

37from lsst.dax.apdb import ( 

38 Apdb, 

39 ApdbConfig, 

40 ApdbReplica, 

41 ApdbTables, 

42 ApdbUpdateRecord, 

43 IncompatibleVersionError, 

44 ReplicaChunk, 

45) 

46from lsst.dax.apdb.pixelization import Pixelization 

47from lsst.dax.apdb.sql import ApdbSql, ApdbSqlConfig 

48from lsst.dax.apdb.tests import ApdbSchemaUpdateTest, ApdbTest 

49 

50try: 

51 import testing.postgresql 

52except ImportError: 

53 testing = None 

54 

55TEST_SCHEMA = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema-apdb.yaml") 

56TEST_SCHEMA_SSO = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema-sso.yaml") 

57# Schema that uses `datetime` for timestamps and combines APDB and SSP. 

58TEST_SCHEMA_DT = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema-datetime.yaml") 

59 

60 

61class ApdbSQLTest(ApdbTest): 

62 """A common base class for SQL APDB tests.""" 

63 

64 dia_object_index: str 

65 

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

67 """Create database and return its config.""" 

68 kw = { 

69 "schema_file": TEST_SCHEMA, 

70 "ss_schema_file": TEST_SCHEMA_SSO, 

71 "dia_object_index": self.dia_object_index, 

72 "enable_replica": self.enable_replica, 

73 } 

74 kw.update(kwargs) 

75 return ApdbSql.init_database(**kw) # type: ignore[arg-type] 

76 

77 def getDiaObjects_table(self) -> ApdbTables: 

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

79 return ApdbTables.DiaObject 

80 

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

82 """Return pixelization used by implementation.""" 

83 assert isinstance(config, ApdbSqlConfig), "Only expect ApdbSqlConfig here" 

84 return Pixelization("htm", config.pixelization.htm_level, config.pixelization.htm_max_ranges) 

85 

86 def test_connection_timeout(self) -> None: 

87 """Test that setting connection timeout does not break things.""" 

88 config = self.make_instance() 

89 assert isinstance(config, ApdbSqlConfig), "Only expect ApdbSqlConfig here" 

90 config.connection_config.connection_timeout = 60.0 

91 with closing(Apdb.from_config(config)): 

92 pass 

93 

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

95 # Docstring inherited. 

96 assert isinstance(apdb, ApdbSql), "Expecting ApdbSql instance" 

97 apdb._storeUpdateRecords(records, chunk, store_chunk=True) 

98 

99 def _count_after_reset_dedup(self, count_before: int) -> int: 

100 return count_before 

101 

102 

103class ApdbSQLiteTestCase(ApdbSQLTest, unittest.TestCase): 

104 """A test case for ApdbSql class using SQLite backend.""" 

105 

106 fsrc_requires_id_list = True 

107 dia_object_index = "baseline" 

108 use_mjd = True 

109 

110 def setUp(self) -> None: 

111 self.tempdir = tempfile.mkdtemp() 

112 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

113 

114 def tearDown(self) -> None: 

115 shutil.rmtree(self.tempdir, ignore_errors=True) 

116 

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

118 return super().make_instance(db_url=self.db_url, **kwargs) 

119 

120 

121class ApdbSQLiteDatetimeTestCase(ApdbSQLTest, unittest.TestCase): 

122 """A test case for ApdbSql class using SQLite backend and schema with 

123 timestamps in TAI MJD. 

124 """ 

125 

126 fsrc_requires_id_list = True 

127 dia_object_index = "baseline" 

128 use_mjd = False 

129 

130 def setUp(self) -> None: 

131 self.tempdir = tempfile.mkdtemp() 

132 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

133 

134 def tearDown(self) -> None: 

135 shutil.rmtree(self.tempdir, ignore_errors=True) 

136 

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

138 if "schema_file" in kwargs: 

139 return super().make_instance(db_url=self.db_url, **kwargs) 

140 else: 

141 return super().make_instance(db_url=self.db_url, schema_file=TEST_SCHEMA_DT, **kwargs) 

142 

143 

144class ApdbSQLiteTestCaseLastObject(ApdbSQLiteTestCase): 

145 """A test case for ApdbSql class using SQLite backend and DiaObjectLast 

146 table. 

147 """ 

148 

149 dia_object_index = "last_object_table" 

150 

151 def getDiaObjects_table(self) -> ApdbTables: 

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

153 return ApdbTables.DiaObjectLast 

154 

155 

156class ApdbSQLiteTestCasePixIdIovIndex(ApdbSQLiteTestCase): 

157 """A test case for ApdbSql class using SQLite backend with pix_id_iov 

158 indexing. 

159 """ 

160 

161 dia_object_index = "pix_id_iov" 

162 

163 

164class ApdbSQLiteTestCaseReplica(ApdbSQLiteTestCase): 

165 """Test case for ApdbSql class using SQLite backend with replica tables.""" 

166 

167 enable_replica = True 

168 meta_row_count = 4 

169 

170 

171@unittest.skipUnless(testing is not None, "testing.postgresql module not found") 

172class ApdbPostgresTestCase(ApdbSQLTest, unittest.TestCase): 

173 """A test case for ApdbSql class using Postgres backend.""" 

174 

175 fsrc_requires_id_list = True 

176 dia_object_index = "last_object_table" 

177 postgresql: Any 

178 enable_replica = True 

179 meta_row_count = 4 

180 timestamp_type_name = "datetime64[ns]" 

181 

182 @classmethod 

183 def setUpClass(cls) -> None: 

184 # Create the postgres test server. 

185 cls.postgresql = testing.postgresql.PostgresqlFactory(cache_initialized_db=True) 

186 super().setUpClass() 

187 

188 @classmethod 

189 def tearDownClass(cls) -> None: 

190 # Clean up any lingering SQLAlchemy engines/connections 

191 # so they're closed before we shut down the server. 

192 gc.collect() 

193 cls.postgresql.clear_cache() 

194 super().tearDownClass() 

195 

196 def setUp(self) -> None: 

197 self.server = self.postgresql() 

198 

199 def tearDown(self) -> None: 

200 self.server = self.postgresql() 

201 

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

203 return super().make_instance(db_url=self.server.url(), **kwargs) 

204 

205 def getDiaObjects_table(self) -> ApdbTables: 

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

207 return ApdbTables.DiaObjectLast 

208 

209 

210@unittest.skipUnless(testing is not None, "testing.postgresql module not found") 

211class ApdbPostgresNamespaceTestCase(ApdbPostgresTestCase): 

212 """A test case for ApdbSql class using Postgres backend with schema name""" 

213 

214 # use mixed case to trigger quoting 

215 namespace = "ApdbSchema" 

216 

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

218 """Make config class instance used in all tests.""" 

219 return super().make_instance(namespace=self.namespace, **kwargs) 

220 

221 

222class ApdbSchemaUpdateSQLiteTestCase(ApdbSchemaUpdateTest, unittest.TestCase): 

223 """A test case for schema updates using SQLite backend.""" 

224 

225 def setUp(self) -> None: 

226 self.tempdir = tempfile.mkdtemp() 

227 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

228 

229 def tearDown(self) -> None: 

230 shutil.rmtree(self.tempdir, ignore_errors=True) 

231 

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

233 """Make config class instance used in all tests.""" 

234 kw = { 

235 "db_url": self.db_url, 

236 "schema_file": TEST_SCHEMA, 

237 "ss_schema_file": TEST_SCHEMA_SSO, 

238 } 

239 kw.update(kwargs) 

240 return ApdbSql.init_database(**kw) # type: ignore[arg-type] 

241 

242 def add_column(self, apdb: Apdb, table: ApdbTables, column: felis.datamodel.Column) -> None: 

243 # SQLAlchemy does not have API for adding columns, have to do it 

244 # manually. 

245 assert isinstance(apdb, ApdbSql) 

246 

247 # Few selected types that we support. 

248 type_names = { 

249 felis.datamodel.DataType.int: "INT", 

250 felis.datamodel.DataType.float: "FLOAT", 

251 felis.datamodel.DataType.double: "DOUBLE", 

252 felis.datamodel.DataType.string: "TEXT", 

253 } 

254 

255 table_name = table.table_name() 

256 datatype = type_names[column.datatype] 

257 not_null = "" if column.nullable else "NOT NULL" 

258 value = "" if column.value is None else column.value 

259 if isinstance(value, str): 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true

260 value = f"'{value}'" 

261 query = f'ALTER TABLE "{table_name}" ADD COLUMN "{column.name}" {datatype} {not_null}' 

262 if value: 262 ↛ 264line 262 didn't jump to line 264 because the condition on line 262 was always true

263 query += f" DEFAULT {value}" 

264 with apdb._engine.begin() as conn: 

265 conn.execute(sqlalchemy.text(query)) 

266 

267 

268class ApdbSQLiteFromUriTestCase(unittest.TestCase): 

269 """A test case for instantiating ApdbSql via URI.""" 

270 

271 def setUp(self) -> None: 

272 self.tempdir = tempfile.mkdtemp() 

273 self.addCleanup(shutil.rmtree, self.tempdir, ignore_errors=True) 

274 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

275 config = ApdbSql.init_database( 

276 db_url=self.db_url, schema_file=TEST_SCHEMA, ss_schema_file=TEST_SCHEMA_SSO 

277 ) 

278 # TODO: This will need update when we switch to pydantic configs. 

279 self.config_path = os.path.join(self.tempdir, "apdb-config.yaml") 

280 config.save(self.config_path) 

281 self.bad_config_path = os.path.join(self.tempdir, "not-config.yaml") 

282 self.index_path = os.path.join(self.tempdir, "apdb-index.yaml") 

283 with open(self.index_path, "w") as index_file: 

284 print(f'label1: "{self.config_path}"', file=index_file) 

285 print(f'"label2/pex_config": "{self.config_path}"', file=index_file) 

286 print(f'bad-label: "{self.bad_config_path}"', file=index_file) 

287 # File with incorrect format. 

288 self.bad_index_path = os.path.join(self.tempdir, "apdb-index-bad.yaml") 

289 with open(self.bad_index_path, "w") as index_file: 

290 print(f'label1: ["{self.config_path}"]', file=index_file) 

291 self.missing_index_path = os.path.join(self.tempdir, "no-apdb-index.yaml") 

292 

293 def test_make_apdb_from_path(self) -> None: 

294 """Check that we can make APDB instance from config URI.""" 

295 with closing(Apdb.from_uri(self.config_path)): 

296 with self.assertRaises(FileNotFoundError): 

297 Apdb.from_uri(self.bad_config_path) 

298 

299 def test_make_apdb_from_labels(self) -> None: 

300 """Check that we can make APDB instance from config URI.""" 

301 # Replace DAX_APDB_INDEX_URI value 

302 new_env = {"DAX_APDB_INDEX_URI": self.index_path} 

303 with patch.dict(os.environ, new_env, clear=True): 

304 with closing(Apdb.from_uri("label:label1")): 

305 pass 

306 with closing(Apdb.from_uri("label:label2")): 

307 pass 

308 # Label does not exist. 

309 with self.assertRaises(ValueError): 

310 Apdb.from_uri("label:not-a-label") 

311 # Label exists but points to a missing config. 

312 with self.assertRaises(FileNotFoundError): 

313 Apdb.from_uri("label:bad-label") 

314 

315 def test_make_apdb_bad_index(self) -> None: 

316 """Check what happens when DAX_APDB_INDEX_URI is broken.""" 

317 # envvar is set but empty. 

318 new_env = {"DAX_APDB_INDEX_URI": ""} 

319 with patch.dict(os.environ, new_env, clear=True): 

320 with self.assertRaises(RuntimeError): 

321 Apdb.from_uri("label:label") 

322 

323 # envvar is set to something non-existing. 

324 new_env = {"DAX_APDB_INDEX_URI": self.missing_index_path} 

325 with patch.dict(os.environ, new_env, clear=True): 

326 with self.assertRaises(FileNotFoundError): 

327 Apdb.from_uri("label:label") 

328 

329 # envvar points to an incorrect file. 

330 new_env = {"DAX_APDB_INDEX_URI": self.bad_index_path} 

331 with patch.dict(os.environ, new_env, clear=True): 

332 with self.assertRaises(TypeError): 

333 Apdb.from_uri("label:label") 

334 

335 def test_remove_database_file(self) -> None: 

336 """Check that SQLite does not try to recreate database after it was 

337 removed, but config persisted. 

338 """ 

339 os.unlink(f"{self.tempdir}/apdb.sqlite3") 

340 with self.assertRaisesRegex(sqlalchemy.exc.OperationalError, "unable to open database file"): 

341 Apdb.from_uri(self.config_path) 

342 

343 def test_make_apdb_replica(self) -> None: 

344 """Check that we can make ApdbReplica instance from config URI.""" 

345 with closing(ApdbReplica.from_uri(self.config_path)): 

346 pass 

347 with self.assertRaises(FileNotFoundError): 

348 Apdb.from_uri(self.bad_config_path) 

349 

350 

351class ApdbSQLiteVersionCheck(unittest.TestCase): 

352 """A test case to verify that version check happens before reading 

353 frozen configuration. 

354 """ 

355 

356 def setUp(self) -> None: 

357 self.tempdir = tempfile.mkdtemp() 

358 self.addCleanup(shutil.rmtree, self.tempdir, ignore_errors=True) 

359 self.db_url = f"sqlite:///{self.tempdir}/apdb.sqlite3" 

360 self.config = ApdbSql.init_database( 

361 db_url=self.db_url, schema_file=TEST_SCHEMA, ss_schema_file=TEST_SCHEMA_SSO 

362 ) 

363 

364 def test_version_check(self) -> None: 

365 """Test that version check happens before reading config.""" 

366 with closing(Apdb.from_config(self.config)) as apdb: 

367 assert isinstance(apdb, ApdbSql) 

368 

369 # Store incompatible version. 

370 apdb.metadata.set(ApdbSql.metadataSchemaVersionKey, "99.0.0", force=True) 

371 

372 # Overwrite frozen config with something that will break. 

373 apdb.metadata.set(ApdbSql.metadataConfigKey, '{"not_a_config_key": 0}', force=True) 

374 

375 # Try again. 

376 with self.assertRaises(IncompatibleVersionError): 

377 Apdb.from_config(self.config) 

378 

379 

380class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase): 

381 """Run file leak tests.""" 

382 

383 

384def setup_module(module: Any) -> None: 

385 """Configure pytest.""" 

386 lsst.utils.tests.init() 

387 

388 

389if __name__ == "__main__": 

390 lsst.utils.tests.init() 

391 unittest.main()