Coverage for tests/test_apdbCassandra.py: 36%

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 `ApdbCassandra` class. 

23 

24Notes 

25----- 

26For now this test can only run against actual Cassandra cluster, to specify 

27cluster location use ``DAX_APDB_TEST_CASSANDRA_CLUSTER`` environment variable, 

28e.g.: 

29 

30 export DAX_APDB_TEST_CASSANDRA_CLUSTER=cassandra.example.com 

31 pytest tests/test_apdbCassandra.py 

32 

33Individual tests create and destroy unique keyspaces in the cluster, there is 

34no need to pre-create a keyspace with predefined name. 

35""" 

36 

37import os 

38import unittest 

39from typing import Any 

40 

41import astropy.time 

42import felis.datamodel 

43 

44import lsst.utils.tests 

45from lsst.dax.apdb import ( 

46 Apdb, 

47 ApdbConfig, 

48 ApdbTables, 

49 ApdbUpdateRecord, 

50 IncompatibleVersionError, 

51 ReplicaChunk, 

52) 

53from lsst.dax.apdb.cassandra import ApdbCassandra, ApdbCassandraConfig 

54from lsst.dax.apdb.cassandra.apdbCassandraSchema import ExtraTables 

55from lsst.dax.apdb.cassandra.connectionContext import ConnectionContext 

56from lsst.dax.apdb.cassandra.sessionFactory import SessionFactory 

57from lsst.dax.apdb.pixelization import Pixelization 

58from lsst.dax.apdb.tests import ApdbSchemaUpdateTest, ApdbTest, cassandra_mixin 

59from lsst.dax.apdb.tests.data_factory import makeObjectCatalog 

60from lsst.dax.apdb.tests.utils import modified_environment 

61 

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

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

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

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

66 

67 

68class ApdbCassandraMixin(cassandra_mixin.ApdbCassandraMixin): 

69 """Mixin class which defines common methods for unit tests.""" 

70 

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

72 """Return pixelization used by implementation.""" 

73 assert isinstance(config, ApdbCassandraConfig), "Only expect ApdbCassandraConfig here" 

74 return Pixelization( 

75 config.partitioning.part_pixelization, 

76 config.partitioning.part_pix_level, 

77 config.partitioning.part_pix_max_ranges, 

78 ) 

79 

80 

81class ApdbCassandraTestCase(ApdbCassandraMixin, ApdbTest, unittest.TestCase): 

82 """A test case for ApdbCassandra class""" 

83 

84 time_partition_tables = False 

85 time_partition_start: str | None = None 

86 time_partition_end: str | None = None 

87 extra_chunk_columns = 2 

88 

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

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

91 kw: dict[str, Any] = { 

92 "hosts": (self.cluster_host,), 

93 "keyspace": self.keyspace, 

94 "schema_file": TEST_SCHEMA, 

95 "ss_schema_file": TEST_SCHEMA_SSO, 

96 "time_partition_tables": self.time_partition_tables, 

97 "enable_replica": self.enable_replica, 

98 } 

99 if self.time_partition_start: 

100 kw["time_partition_start"] = self.time_partition_start 

101 if self.time_partition_end: 

102 kw["time_partition_end"] = self.time_partition_end 

103 kw.update(kwargs) 

104 return ApdbCassandra.init_database(**kw) 

105 

106 def getDiaObjects_table(self) -> ApdbTables: 

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

108 return ApdbTables.DiaObjectLast 

109 

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

111 # Docstring inherited. 

112 assert isinstance(apdb, ApdbCassandra), "Expecting ApdbCassandra instance" 

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

114 

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

116 return 0 

117 

118 

119class ApdbCassandraPerMonthTestCase(ApdbCassandraTestCase): 

120 """A test case for ApdbCassandra class with per-month tables.""" 

121 

122 time_partition_tables = True 

123 time_partition_start = "2020-06-01T00:00:00" 

124 time_partition_end = "2021-06-01T00:00:00" 

125 meta_row_count = 4 

126 

127 def test_store_partition_range(self) -> None: 

128 """Test that writing to non-existing partition raises an error.""" 

129 config = self.make_instance() 

130 apdb = Apdb.from_config(config) 

131 

132 region = self.make_region() 

133 

134 # Visit time is beyond time_partition_end. 

135 visit_time = astropy.time.Time("2022-01-01", format="isot", scale="tai") 

136 catalog = makeObjectCatalog(region, 100) 

137 with self.assertRaisesRegex(ValueError, "time partitions that do not yet exist"): 

138 apdb.store(visit_time, catalog) 

139 

140 # Writing to last partition makes a warning. 

141 visit_time = astropy.time.Time("2021-06-01", format="isot", scale="tai") 

142 catalog = makeObjectCatalog(region, 100) 

143 with self.assertWarnsRegex(UserWarning, "Writing into the last temporal partition"): 

144 apdb.store(visit_time, catalog) 

145 

146 

147class ApdbCassandraTestCaseReplica(ApdbCassandraTestCase): 

148 """A test case with enabled replica tables.""" 

149 

150 enable_replica = True 

151 meta_row_count = 4 

152 

153 

154class ApdbCassandraTestCaseDatetimeReplica(ApdbCassandraTestCaseReplica): 

155 """A test case with datetime timestamps.""" 

156 

157 use_mjd = False 

158 

159 def setUp(self) -> None: 

160 super().setUp() 

161 # Schema for datetime case is also missing a validityTime column in 

162 # DiaObjectLast table. 

163 self.table_column_count = dict(self.table_column_count) 

164 self.table_column_count[ApdbTables.DiaObjectLast] = 5 

165 

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

167 if "schema_file" in kwargs: 

168 return super().make_instance(**kwargs) 

169 else: 

170 return super().make_instance(schema_file=TEST_SCHEMA_DT, **kwargs) 

171 

172 

173class ApdbSchemaUpdateCassandraTestCase(ApdbCassandraMixin, ApdbSchemaUpdateTest, unittest.TestCase): 

174 """A test case for schema updates using Cassandra backend.""" 

175 

176 time_partition_tables = False 

177 

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

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

180 kw = { 

181 "hosts": (self.cluster_host,), 

182 "keyspace": self.keyspace, 

183 "schema_file": TEST_SCHEMA, 

184 "ss_schema_file": TEST_SCHEMA_SSO, 

185 "time_partition_tables": self.time_partition_tables, 

186 } 

187 kw.update(kwargs) 

188 return ApdbCassandra.init_database(**kw) # type: ignore[arg-type] 

189 

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

191 assert isinstance(apdb, ApdbCassandra) 

192 

193 # Few selected types that we support. 

194 type_names = { 

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

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

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

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

199 } 

200 

201 context = apdb._context 

202 # Need to add to both regular and replica tables. 

203 table_names = [table.table_name(), ExtraTables.replica_chunk_tables(True)[table].table_name()] 

204 datatype = type_names[column.datatype] 

205 for table_name in table_names: 

206 query = f'ALTER TABLE "{context.config.keyspace}"."{table_name}" ADD "{column.name}" {datatype}' 

207 context.session.execute(query) 

208 

209 

210class ApdbCassandraVersionCheck(cassandra_mixin.ApdbCassandraMixin, unittest.TestCase): 

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

212 frozen configuration. 

213 """ 

214 

215 def setUp(self) -> None: 

216 super().setUp() 

217 

218 self.config = ApdbCassandra.init_database( 

219 hosts=(self.cluster_host,), 

220 keyspace=self.keyspace, 

221 schema_file=TEST_SCHEMA, 

222 ss_schema_file=TEST_SCHEMA_SSO, 

223 time_partition_tables=False, 

224 ) 

225 

226 def test_version_check(self) -> None: 

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

228 apdb = Apdb.from_config(self.config) 

229 assert isinstance(apdb, ApdbCassandra) 

230 

231 # Store incompatible version. 

232 apdb.metadata.set(ConnectionContext.metadataSchemaVersionKey, "99.0.0", force=True) 

233 

234 # Overwrite frozen config with something that will break. 

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

236 

237 # Try again. 

238 with self.assertRaises(IncompatibleVersionError): 

239 # Need to call some actual method to initiate connection. 

240 Apdb.from_config(self.config).metadata.items() 

241 

242 

243_DB_AUTH_JSON = """\ 

244[{ 

245 "url": "cassandra://user1000@node1.slac.stanford.edu:9042/", 

246 "username": "user01", 

247 "password": "pass01" 

248}, { 

249 "url": "cassandra://node2.slac.stanford.edu:9042/", 

250 "username": "user02", 

251 "password": "pass02" 

252}, { 

253 "url": "cassandra://node1.slac.stanford.edu:9042/apdb_dev", 

254 "username": "user03", 

255 "password": "pass03" 

256}, { 

257 "url": "cassandra://user2000@test_cluster/", 

258 "username": "user04", 

259 "password": "pass04" 

260}, { 

261 "url": "cassandra://test_cluster/", 

262 "username": "user05", 

263 "password": "pass05" 

264}, { 

265 "url": "cassandra://pp_apdb_prod_cluster/", 

266 "username": "user06", 

267 "password": "pass06" 

268}, { 

269 "url": "cassandra://pp_apdb_dev_cluster/", 

270 "username": "user07", 

271 "password": "pass07" 

272}] 

273""" 

274 

275 

276class ApdbCassandraDbAuthTest(unittest.TestCase): 

277 """A test case for extracting credentials from db-auth.yaml.""" 

278 

279 def _make_config(self) -> ApdbCassandraConfig: 

280 config = ApdbCassandraConfig( 

281 contact_points=("node1.slac.stanford.edu", "node2.slac.stanford.edu"), 

282 keyspace="apdb", 

283 ) 

284 return config 

285 

286 @unittest.skipIf(not cassandra_mixin.CASSANDRA_IMPORTED, "cassandra_driver cannot be imported") 

287 def test_dbauth(self) -> None: 

288 """Check credentials access.""" 

289 with modified_environment(LSST_DB_AUTH_CREDENTIALS=_DB_AUTH_JSON): 

290 config = self._make_config() 

291 

292 factory = SessionFactory(config) 

293 

294 # Should match second entry. 

295 auth = factory._make_auth_provider() 

296 assert auth is not None 

297 self.assertEqual(auth.username, "user02") 

298 

299 config.keyspace = "apdb_dev" 

300 # Should match third entry. 

301 auth = factory._make_auth_provider() 

302 assert auth is not None 

303 self.assertEqual(auth.username, "user03") 

304 

305 config.connection_config.username = "user1000" 

306 # Should match first entry, returns original user name. 

307 auth = factory._make_auth_provider() 

308 assert auth is not None 

309 self.assertEqual(auth.username, "user1000") 

310 self.assertEqual(auth.password, "pass01") 

311 

312 config.connection_config.username = "" 

313 config.connection_config.dbauth_alias = "test_cluster" 

314 # Should match fifth entry. 

315 auth = factory._make_auth_provider() 

316 assert auth is not None 

317 self.assertEqual(auth.username, "user05") 

318 

319 config.connection_config.username = "user2000" 

320 config.connection_config.dbauth_alias = "test_cluster" 

321 # Should match fourth entry. 

322 auth = factory._make_auth_provider() 

323 assert auth is not None 

324 self.assertEqual(auth.username, "user2000") 

325 self.assertEqual(auth.password, "pass04") 

326 

327 # Default value of alias is based on username and keyspace. 

328 config.connection_config.dbauth_alias = "" 

329 config.contact_points = ("127.0.0.1",) 

330 config.keyspace = "apdb" 

331 config.connection_config.username = "apdb-prod" 

332 auth = factory._make_auth_provider() 

333 self.assertIsNone(auth) 

334 

335 config.keyspace = "pp_apdb_cam" 

336 auth = factory._make_auth_provider() 

337 assert auth is not None 

338 self.assertEqual(auth.password, "pass06") 

339 

340 config.connection_config.username = "apdb" 

341 config.keyspace = "pp_apdb_cam_dev" 

342 auth = factory._make_auth_provider() 

343 assert auth is not None 

344 self.assertEqual(auth.password, "pass07") 

345 

346 

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

348 """Run file leak tests.""" 

349 

350 

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

352 """Configure pytest.""" 

353 lsst.utils.tests.init() 

354 

355 

356if __name__ == "__main__": 

357 lsst.utils.tests.init() 

358 unittest.main()