asyncprocess/tests/db_read_tests.py

0001 # Description: Unit tests for the shape and cost checks and the SQL of the async db_read requests.
0002 # The tables are described the way PanDAWMS/panda-database schema/oracle has them; no database is used
0003 import datetime
0004 import unittest
0005 from typing import Any
0006 from unittest import mock
0007
0008 from pandaserver.asyncprocess.db_read import schema_info
0009 from pandaserver.asyncprocess.db_read.compose import compose_statement, escape_prefix
0010 from pandaserver.asyncprocess.db_read.request_shape import (
0011     MAX_KEYS,
0012     FromQuery,
0013     Query,
0014     QueryRejected,
0015     TablePlan,
0016     check_query,
0017     count_touched_partitions,
0018     parse_query,
0019 )
0020 from pandaserver.asyncprocess.db_read.schema_info import (
0021     KIND_LOB,
0022     KIND_NUMBER,
0023     KIND_OTHER,
0024     KIND_TEXT,
0025     KIND_TIME,
0026     STATEMENT_COMMENT,
0027     IndexInfo,
0028     PartitionInfo,
0029     TableInfo,
0030     TableNotVisible,
0031     get_table_info,
0032     read_table_info,
0033 )
0034 from pandaserver.taskbuffer.WrappedCursor import convert_query_in_printf_format
0035
0036 DAY = datetime.timedelta(days=1)
0037 D0 = datetime.datetime(2026, 10, 1)
0038
0039 JOB_COLUMNS = {
0040     "pandaid": KIND_NUMBER,
0041     "jeditaskid": KIND_NUMBER,
0042     "jobstatus": KIND_TEXT,
0043     "computingsite": KIND_TEXT,
0044     "modificationtime": KIND_TIME,
0045 }
0046
0047 # not partitioned, so every index is one structure
0048 JOBS_ACTIVE4 = TableInfo(
0049     "ATLAS_PANDA",
0050     "jobsActive4",
0051     dict(JOB_COLUMNS),
0052     (IndexInfo("pandaid", False), IndexInfo("jeditaskid", False), IndexInfo("jobstatus", False), IndexInfo("computingsite", False)),
0053     None,
0054 )
0055 JOBS_DEFINED4 = TableInfo("ATLAS_PANDA", "jobsDefined4", dict(JOB_COLUMNS), (IndexInfo("pandaid", False),), None)
0056 # daily partitions on modificationTime: the range partition it is anchored on, then four
0057 # interval partitions, each starting where the one before ends as Oracle tells it
0058 JOBS_ARCHIVED4 = TableInfo(
0059     "ATLAS_PANDA",
0060     "jobsArchived4",
0061     dict(JOB_COLUMNS),
0062     (IndexInfo("pandaid", True), IndexInfo("jeditaskid", True)),
0063     PartitionInfo("modificationtime", ((None, D0), (D0, D0 + DAY), (D0 + DAY, D0 + 2 * DAY), (D0 + 2 * DAY, D0 + 3 * DAY), (D0 + 3 * DAY, D0 + 4 * DAY))),
0064 )
0065 # partitions of 500000 tasks on jediTaskID, with its primary key LOCAL
0066 JEDI_TASKS = TableInfo(
0067     "ATLAS_PANDA",
0068     "JEDI_Tasks",
0069     {"jeditaskid": KIND_NUMBER, "status": KIND_TEXT, "modificationtime": KIND_TIME, "splitrule": KIND_TEXT},
0070     (IndexInfo("jeditaskid", True), IndexInfo("status", True)),
0071     PartitionInfo("jeditaskid", ((None, 500000), (500000, 1000000), (1000000, 1500000))),
0072 )
0073 # a few hundred rows, with a primary key that starts with app
0074 CONFIG = TableInfo(
0075     "ATLAS_PANDA",
0076     "CONFIG",
0077     {"app": KIND_TEXT, "component": KIND_TEXT, "key": KIND_TEXT, "value": KIND_TEXT, "vo": KIND_TEXT},
0078     (IndexInfo("app", False),),
0079     None,
0080 )
0081 JOB_PARAMS = TableInfo(
0082     "ATLAS_PANDA",
0083     "jobParamsTable",
0084     {"pandaid": KIND_NUMBER, "modificationtime": KIND_TIME, "jobparameters": KIND_LOB},
0085     (IndexInfo("pandaid", True),),
0086     PartitionInfo("modificationtime", ((None, D0), (D0, D0 + DAY))),
0087 )
0088
0089 INFOS = {info.table_name: info for info in (JOBS_ACTIVE4, JOBS_DEFINED4, JOBS_ARCHIVED4, JEDI_TASKS, CONFIG, JOB_PARAMS)}
0090
0091
0092 def parse(raw: dict[str, Any], position: int = 0) -> Query:
0093     return parse_query(raw, position)
0094
0095
0096 def check(raw: dict[str, Any], budget: int = 1000, from_sizes: dict[FromQuery, int] | None = None, position: int = 0) -> list[TablePlan]:
0097     return check_query(parse(raw, position), INFOS, from_sizes or {}, budget)
0098
0099
0100 class TestParseQuery(unittest.TestCase):
0101     def assertRejected(self, rule: str, raw: Any, position: int = 0) -> QueryRejected:
0102         with self.assertRaises(QueryRejected) as context:
0103             parse_query(raw, position)
0104         self.assertEqual(context.exception.rule, rule, context.exception.message)
0105         return context.exception
0106
0107     def test_full_query(self) -> None:
0108         query = parse(
0109             {
0110                 "tables": ["JOBSACTIVE4", "jobsArchived4"],
0111                 "where": [["PandaID", "IN", [1, 2]], ["modificationTime", ">=", {"__datetime__": "2026-10-01T00:00:00"}], ["jobStatus", "is  not null"]],
0112                 "group_by": ["computingSite", "day(modificationTime)"],
0113                 "aggregates": [["count", "*"], ["MAX", "modificationTime"]],
0114                 "order_by": [["count(*)", "desc"], "computingsite"],
0115                 "limit": 10,
0116             }
0117         )
0118         self.assertEqual([table.name for table in query.tables], ["jobsActive4", "jobsArchived4"])
0119         self.assertEqual(
0120             [(c.column, c.op, c.value) for c in query.where],
0121             [("pandaid", "in", [1, 2]), ("modificationtime", ">=", "2026-10-01T00:00:00"), ("jobstatus", "is not null", None)],
0122         )
0123         self.assertEqual([key.name for key in query.group_by], ["computingsite", "day(modificationtime)"])
0124         self.assertEqual([aggregate.name for aggregate in query.aggregates], ["count(*)", "max(modificationtime)"])
0125         self.assertEqual([(key.name, key.descending) for key in query.order_by], [("count(*)", True), ("computingsite", False)])
0126         self.assertEqual(query.limit, 10)
0127
0128     def test_defaults(self) -> None:
0129         query = parse({"tables": ["CONFIG"]})
0130         self.assertEqual((query.where, query.columns, query.limit), ((), None, 1000))
0131
0132     def test_tables(self) -> None:
0133         self.assertRejected("unknown_table", {"tables": ["users"]})
0134         self.assertRejected("unknown_table", {"tables": ["jobs_StatusLog"]})
0135         self.assertRejected("bad_shape", {"tables": []})
0136         self.assertRejected("bad_shape", {"tables": ["jobsActive4", "JOBSACTIVE4"]})
0137         self.assertRejected("bad_shape", {"tables": ["jobsActive4"], "colums": ["pandaid"]})
0138
0139     def test_conditions(self) -> None:
0140         tables = ["jobsActive4"]
0141         self.assertRejected("bad_shape", {"tables": tables, "where": [["pandaid", "like", "1%"]]})
0142         self.assertRejected("bad_shape", {"tables": tables, "where": [["pandaid", "is null", 1]]})
0143         self.assertRejected("bad_shape", {"tables": tables, "where": [["pandaid", "="]]})
0144         self.assertRejected("bad_shape", {"tables": tables, "where": [["pandaid; drop", "=", 1]]})
0145         self.assertRejected("bad_value", {"tables": tables, "where": [["pandaid", "in", []]]})
0146         self.assertRejected("bad_value", {"tables": tables, "where": [["pandaid", "=", True]]})
0147         self.assertRejected("bad_value", {"tables": tables, "where": [["pandaid", "=", None]]})
0148         self.assertRejected("bad_value", {"tables": tables, "where": [["jobstatus", "prefix", ""]]})
0149         rejected = self.assertRejected("too_many_keys", {"tables": tables, "where": [["pandaid", "in", list(range(MAX_KEYS + 1))]]})
0150         self.assertEqual(rejected.detail, {"keys": MAX_KEYS + 1, "max_keys": MAX_KEYS})
0151         self.assertRejected("bad_shape", {"tables": tables, "where": [["pandaid", "=", 1]] * 21})
0152
0153     def test_from_query(self) -> None:
0154         raw = {"tables": ["jobsActive4"], "where": [["jeditaskid", "in", {"from_query": 0, "column": "jediTaskID"}]]}
0155         self.assertEqual(parse(raw, position=1).where[0].value, FromQuery(0, "jeditaskid"))
0156         # only an earlier query
0157         self.assertRejected("bad_shape", raw, position=0)
0158         self.assertRejected("bad_shape", {"tables": ["jobsActive4"], "where": [["jeditaskid", "=", {"from_query": 0, "column": "c"}]]}, position=1)
0159         self.assertRejected("bad_shape", {"tables": ["jobsActive4"], "where": [["jeditaskid", "in", {"from_query": 0}]]}, position=1)
0160
0161     def test_returned(self) -> None:
0162         tables = ["jobsActive4"]
0163         self.assertRejected("bad_shape", {"tables": tables, "columns": ["pandaid"], "group_by": ["jobstatus"]})
0164         self.assertRejected("bad_shape", {"tables": tables, "columns": []})
0165         self.assertRejected("bad_shape", {"tables": tables, "columns": ["pandaid", "PandaID"]})
0166         self.assertRejected("bad_shape", {"tables": tables, "aggregates": [["min", "*"]]})
0167         self.assertRejected("bad_shape", {"tables": tables, "aggregates": [["sum", "pandaid"]]})
0168         self.assertRejected("bad_shape", {"tables": tables, "columns": ["pandaid"], "order_by": ["jobstatus"]})
0169         self.assertRejected("bad_shape", {"tables": tables, "group_by": ["jobstatus"], "order_by": ["pandaid"]})
0170         self.assertRejected("bad_shape", {"tables": tables, "order_by": [["pandaid", "down"]]})
0171         self.assertRejected("bad_shape", {"tables": tables, "limit": 1001})
0172         self.assertRejected("bad_shape", {"tables": tables, "limit": 0})
0173         self.assertRejected("bad_shape", {"tables": tables, "limit": True})
0174
0175
0176 class TestCountTouchedPartitions(unittest.TestCase):
0177     def count(self, where: list[list[Any]], now: datetime.datetime | None = None) -> int:
0178         plans = check({"tables": ["jobsArchived4"], "where": [["pandaid", "=", 1]] + where})
0179         assert JOBS_ARCHIVED4.partitioning is not None
0180         return count_touched_partitions(JOBS_ARCHIVED4.partitioning, plans[0].where, now)
0181
0182     def test_window(self) -> None:
0183         self.assertEqual(self.count([]), 5)
0184         # the partition [D0+1, D0+2) and the ones after it
0185         self.assertEqual(self.count([["modificationtime", ">=", "2026-10-02T00:00:00"]]), 3)
0186         self.assertEqual(self.count([["modificationtime", ">", "2026-10-02T00:00:00"]]), 3)
0187         # an upper bound at the start of a partition leaves it out only when exclusive
0188         self.assertEqual(self.count([["modificationtime", "<", "2026-10-02T00:00:00"]]), 2)
0189         self.assertEqual(self.count([["modificationtime", "<=", "2026-10-02T00:00:00"]]), 3)
0190         self.assertEqual(self.count([["modificationtime", ">=", "2026-10-02T06:00:00"], ["modificationtime", "<", "2026-10-02T18:00:00"]]), 1)
0191         # the narrower of two lower bounds wins
0192         self.assertEqual(self.count([["modificationtime", ">=", "2026-10-01T00:00:00"], ["modificationtime", ">=", "2026-10-04T00:00:00"]]), 1)
0193         # before the oldest bound, only the partition open below
0194         self.assertEqual(self.count([["modificationtime", "<", "2026-09-01T00:00:00"]]), 1)
0195
0196     def test_points(self) -> None:
0197         self.assertEqual(self.count([["modificationtime", "=", "2026-10-02T12:00:00"]]), 1)
0198         self.assertEqual(self.count([["modificationtime", "in", ["2026-10-02T12:00:00", "2026-10-03T12:00:00", "2026-10-03T13:00:00"]]]), 2)
0199         # a point at a bound is in the partition that starts there
0200         self.assertEqual(self.count([["modificationtime", "in", ["2026-10-02T00:00:00"]]]), 1)
0201         # a point outside the window leaves nothing to read, and the count does not go below 1
0202         self.assertEqual(self.count([["modificationtime", "in", ["2026-10-02T12:00:00"]], ["modificationtime", ">=", "2026-10-03T00:00:00"]]), 1)
0203
0204     def test_made_ahead(self) -> None:
0205         # on 2026-10-02, the partitions of 2026-10-03 and 2026-10-04 are made but hold no row
0206         now = D0 + DAY + datetime.timedelta(hours=12)
0207         self.assertEqual(self.count([], now), 3)
0208         self.assertEqual(self.count([["modificationtime", ">=", "2026-10-02T00:00:00"]], now), 1)
0209         # the partition holding now counts from its first second
0210         self.assertEqual(self.count([], D0 + DAY), 3)
0211         # a point ahead is in no partition counted, and the count does not go below 1
0212         self.assertEqual(self.count([["modificationtime", "=", "2026-10-04T12:00:00"]], now), 1)
0213         # a key that is not a time keeps every partition
0214         assert JEDI_TASKS.partitioning is not None
0215         plans = check({"tables": ["JEDI_Tasks"], "where": [["jeditaskid", ">=", 600000], ["status", "=", "running"]]})
0216         self.assertEqual(count_touched_partitions(JEDI_TASKS.partitioning, plans[0].where, now), 2)
0217         # and the cost of a query follows
0218         query = parse({"tables": ["jobsArchived4"], "where": [["pandaid", "in", [1, 2]]]}, 0)
0219         self.assertEqual(check_query(query, INFOS, {}, 1000)[0].probes, 10)
0220         self.assertEqual(check_query(query, INFOS, {}, 1000, now=now)[0].probes, 6)
0221
0222     def test_numbers(self) -> None:
0223         assert JEDI_TASKS.partitioning is not None
0224         plans = check({"tables": ["JEDI_Tasks"], "where": [["jeditaskid", "in", [10, 499999, 500000]]]})
0225         self.assertEqual(count_touched_partitions(JEDI_TASKS.partitioning, plans[0].where), 2)
0226         plans = check({"tables": ["JEDI_Tasks"], "where": [["jeditaskid", ">=", 600000], ["status", "=", "running"]]})
0227         self.assertEqual(count_touched_partitions(JEDI_TASKS.partitioning, plans[0].where), 2)
0228
0229
0230 class TestCheckQuery(unittest.TestCase):
0231     def assertRejected(self, rule: str, raw: dict[str, Any], **kwargs: Any) -> QueryRejected:
0232         with self.assertRaises(QueryRejected) as context:
0233             check(raw, **kwargs)
0234         self.assertEqual(context.exception.rule, rule, context.exception.message)
0235         return context.exception
0236
0237     def test_job_tables(self) -> None:
0238         plans = check({"tables": ["jobsDefined4", "jobsActive4", "jobsArchived4"], "where": [["pandaid", "in", [1, 2, 3]]]})
0239         # a LOCAL index is probed in each of the five partitions
0240         self.assertEqual([plan.probes for plan in plans], [3, 3, 15])
0241         plans = check({"tables": ["jobsArchived4"], "where": [["pandaid", "in", [1, 2, 3]], ["modificationtime", ">=", "2026-10-03T00:00:00"]]})
0242         self.assertEqual(plans[0].probes, 6)
0243
0244     def test_cheapest_index(self) -> None:
0245         plans = check({"tables": ["jobsActive4"], "where": [["jobstatus", "=", "running"], ["pandaid", "in", [1, 2]]]})
0246         self.assertEqual(plans[0].probes, 1)
0247
0248     def test_partition_key_index(self) -> None:
0249         # each task is in one partition
0250         plans = check({"tables": ["JEDI_Tasks"], "where": [["jeditaskid", "in", [10, 600000, 1200000]]]})
0251         self.assertEqual(plans[0].probes, 3)
0252         # status is not the partition key, so its LOCAL index is probed in each partition
0253         plans = check({"tables": ["JEDI_Tasks"], "where": [["status", "=", "running"]]})
0254         self.assertEqual(plans[0].probes, 3)
0255
0256     def test_budget(self) -> None:
0257         rejected = self.assertRejected("probe_budget", {"tables": ["jobsActive4", "jobsArchived4"], "where": [["pandaid", "in", [1, 2, 3]]]}, budget=17)
0258         self.assertEqual(rejected.detail, {"probes": {"jobsActive4": 3, "jobsArchived4": 15}, "budget": 17})
0259         check({"tables": ["jobsActive4", "jobsArchived4"], "where": [["pandaid", "in", [1, 2, 3]]]}, budget=18)
0260
0261     def test_no_index_key(self) -> None:
0262         # a prefix is no key, though computingsite leads an index
0263         self.assertRejected("no_index_key", {"tables": ["jobsActive4"], "where": [["computingsite", "prefix", "CERN"]]})
0264         self.assertRejected("no_index_key", {"tables": ["jobsActive4"], "where": [["pandaid", ">", 1]]})
0265         # a window on the partition key alone is not enough
0266         self.assertRejected("no_index_key", {"tables": ["jobsArchived4"], "where": [["modificationtime", ">=", "2026-10-03T00:00:00"]]})
0267         # however small the table
0268         self.assertRejected("no_index_key", {"tables": ["CONFIG"], "where": [["component", "=", "jedi"], ["key", "=", "x"]]})
0269         self.assertRejected("no_index_key", {"tables": ["CONFIG"]})
0270         plans = check({"tables": ["CONFIG"], "where": [["app", "in", ["jedi", "pandaserver"]], ["component", "=", "jedi"], ["key", "=", "x"]]})
0271         self.assertEqual(plans[0].probes, 2)
0272
0273     def test_values(self) -> None:
0274         # a text for a number column would make the database convert the column
0275         rejected = self.assertRejected("bad_value", {"tables": ["jobsActive4"], "where": [["pandaid", "=", "1"]]})
0276         self.assertEqual(rejected.detail, {"table": "jobsActive4", "column": "pandaid"})
0277         self.assertRejected("bad_value", {"tables": ["jobsActive4"], "where": [["jobstatus", "=", 1], ["pandaid", "=", 1]]})
0278         self.assertRejected("bad_value", {"tables": ["jobsActive4"], "where": [["pandaid", "=", 1], ["modificationtime", ">", "yesterday"]]})
0279         self.assertRejected("bad_value", {"tables": ["jobsActive4"], "where": [["pandaid", "prefix", "1"]]})
0280         self.assertRejected("bad_value", {"tables": ["jobParamsTable"], "where": [["pandaid", "=", 1], ["jobparameters", "prefix", "x"]]})
0281         plans = check({"tables": ["jobsActive4"], "where": [["pandaid", "=", 1], ["modificationtime", ">", "2026-10-01T02:00:00+02:00"]]})
0282         self.assertEqual(plans[0].where[1].value, datetime.datetime(2026, 10, 1, 0, 0))
0283
0284     def test_columns(self) -> None:
0285         rejected = self.assertRejected("unknown_column", {"tables": ["jobsActive4"], "columns": ["nosuch"], "where": [["pandaid", "=", 1]]})
0286         self.assertEqual(rejected.detail, {"table": "jobsActive4", "column": "nosuch"})
0287         self.assertRejected("unknown_column", {"tables": ["jobsActive4"], "where": [["pandaid", "=", 1]], "order_by": ["nosuch"]})
0288         self.assertRejected("bad_value", {"tables": ["jobsActive4"], "where": [["pandaid", "=", 1]], "group_by": ["day(jobstatus)"]})
0289         self.assertRejected("bad_value", {"tables": ["jobParamsTable"], "where": [["pandaid", "=", 1]], "aggregates": [["max", "jobparameters"]]})
0290
0291     def test_from_query(self) -> None:
0292         reference = FromQuery(0, "pandaid")
0293         raw = {"tables": ["jobsArchived4"], "where": [["pandaid", "in", {"from_query": 0, "column": "pandaid"}]]}
0294         plans = check(raw, from_sizes={reference: 4}, position=1)
0295         self.assertEqual(plans[0].probes, 20)
0296         rejected = self.assertRejected("too_many_keys", raw, from_sizes={reference: MAX_KEYS + 1}, position=1)
0297         self.assertEqual(rejected.detail, {"keys": MAX_KEYS + 1, "max_keys": MAX_KEYS})
0298
0299
0300 class TestCompose(unittest.TestCase):
0301     def compose(self, raw: dict[str, Any], backend: str, from_values: dict[FromQuery, list[Any]] | None = None, position: int = 0) -> Any:
0302         query = parse(raw, position)
0303         from_values = from_values or {}
0304         plans = check_query(query, INFOS, {reference: len(values) for reference, values in from_values.items()}, 1000)
0305         return [compose_statement(plan, query, from_values, backend) for plan in plans]
0306
0307     def test_oracle(self) -> None:
0308         (statement,) = self.compose(
0309             {
0310                 "tables": ["jobsArchived4"],
0311                 "where": [
0312                     ["pandaid", "in", [1, 2]],
0313                     ["modificationtime", ">=", "2026-10-01T00:00:00"],
0314                     ["jobstatus", "!=", "failed"],
0315                     ["computingsite", "is null"],
0316                 ],
0317                 "columns": ["pandaid", "jobstatus"],
0318                 "order_by": [["pandaid", "desc"]],
0319                 "limit": 100,
0320             },
0321             "oracle",
0322         )
0323         self.assertEqual(
0324             statement.sql,
0325             "SELECT pandaid,jobstatus FROM ATLAS_PANDA.jobsArchived4 "
0326             "WHERE pandaid IN (:r0,:r1) AND modificationtime >= :r2 AND jobstatus <> :r3 AND computingsite IS NULL "
0327             "ORDER BY pandaid DESC FETCH FIRST 101 ROWS ONLY" + STATEMENT_COMMENT,
0328         )
0329         self.assertEqual(statement.binds, {":r0": 1, ":r1": 2, ":r2": D0, ":r3": "failed"})
0330         self.assertEqual((statement.names, statement.limit), (("pandaid", "jobstatus"), 100))
0331
0332     def test_grouped(self) -> None:
0333         raw = {
0334             "tables": ["jobsActive4"],
0335             "where": [["jeditaskid", "=", 7]],
0336             "group_by": ["jobstatus", "hour(modificationtime)"],
0337             "aggregates": [["count", "*"], ["min", "pandaid"]],
0338             "order_by": ["hour(modificationtime)"],
0339         }
0340         (oracle,) = self.compose(raw, "oracle")
0341         self.assertEqual(
0342             oracle.sql,
0343             "SELECT jobstatus,TRUNC(modificationtime,'HH24'),COUNT(*),MIN(pandaid) FROM ATLAS_PANDA.jobsActive4 WHERE jeditaskid = :r0 "
0344             "GROUP BY jobstatus,TRUNC(modificationtime,'HH24') ORDER BY TRUNC(modificationtime,'HH24') FETCH FIRST 1001 ROWS ONLY" + STATEMENT_COMMENT,
0345         )
0346         self.assertEqual(oracle.names, ("jobstatus", "hour(modificationtime)", "count(*)", "min(pandaid)"))
0347         (postgres,) = self.compose(raw, "postgres")
0348         self.assertIn("date_trunc('hour',modificationtime)", postgres.sql)
0349         self.assertTrue(postgres.sql.endswith(" LIMIT 1001" + STATEMENT_COMMENT))
0350
0351     def test_postgres_conversion(self) -> None:
0352         # what WrappedCursor makes of the statement: one %s per bind variable, in order, and nothing else touched
0353         (statement,) = self.compose(
0354             {
0355                 "tables": ["jobsActive4"],
0356                 "where": [["pandaid", "in", [1, 2]], ["computingsite", "prefix", "CERN_1%"], ["modificationtime", "<", "2026-10-01T00:00:00"]],
0357                 "group_by": ["day(modificationtime)"],
0358                 "aggregates": [["count", "*"]],
0359             },
0360             "postgres",
0361         )
0362         sql, params = convert_query_in_printf_format(statement.sql, [statement.binds], {})
0363         # the comment comes through the conversion as it is, which a ":" or a "%" in it would not
0364         self.assertEqual(convert_query_in_printf_format("SELECT 1" + STATEMENT_COMMENT, [{}], {})[0], "SELECT 1" + STATEMENT_COMMENT)
0365         self.assertEqual(
0366             sql,
0367             "SELECT date_trunc('day',modificationtime),COUNT(*) FROM ATLAS_PANDA.jobsActive4 "
0368             "WHERE pandaid IN (%s,%s) AND computingsite LIKE %s ESCAPE '!' AND modificationtime < %s "
0369             "GROUP BY date_trunc('day',modificationtime) LIMIT 1001" + STATEMENT_COMMENT,
0370         )
0371         self.assertEqual(params, [[1, 2, "CERN!_1!%%", D0]])
0372
0373     def test_default_columns_leave_out_large_objects(self) -> None:
0374         (statement,) = self.compose({"tables": ["jobParamsTable"], "where": [["pandaid", "=", 1]]}, "oracle")
0375         self.assertEqual(statement.names, ("pandaid", "modificationtime"))
0376         (statement,) = self.compose({"tables": ["jobParamsTable"], "where": [["pandaid", "=", 1]], "columns": ["jobparameters"]}, "oracle")
0377         self.assertEqual(statement.names, ("jobparameters",))
0378
0379     def test_from_query(self) -> None:
0380         raw = {"tables": ["jobsActive4"], "where": [["pandaid", "in", {"from_query": 0, "column": "pandaid"}]]}
0381         (statement,) = self.compose(raw, "oracle", from_values={FromQuery(0, "pandaid"): [5, 6]}, position=1)
0382         self.assertEqual(statement.binds, {":r0": 5, ":r1": 6})
0383         # no values, so no row can match and there is nothing to run
0384         self.assertEqual(self.compose(raw, "oracle", from_values={FromQuery(0, "pandaid"): []}, position=1), [None])
0385
0386     def test_escape_prefix(self) -> None:
0387         self.assertEqual(escape_prefix("a_b%c!"), "a!_b!%c!!%")
0388
0389
0390 def oracle_date(day: str) -> str:
0391     return f"TO_DATE(' {day} 00:00:00', 'SYYYY-MM-DD HH24:MI:SS', 'NLS_CALENDAR=GREGORIAN')"
0392
0393
0394 class FakeDictionary:
0395     # answers the statements of schema_info by the view each reads from, the one after FROM
0396     def __init__(self, answers: dict[str, list[tuple[Any, ...]]]) -> None:
0397         self.answers = answers
0398         self.statements: list[tuple[str, dict[str, Any]]] = []
0399
0400     def __call__(self, sql: str, binds: dict[str, Any]) -> list[tuple[Any, ...]]:
0401         self.statements.append((sql, binds))
0402         for view, rows in self.answers.items():
0403             if f"FROM {view} " in sql:
0404                 return rows
0405         return []
0406
0407
0408 class TestSchemaInfo(unittest.TestCase):
0409     def setUp(self) -> None:
0410         schema_info._cache.clear()
0411
0412     def test_oracle_time_partitions(self) -> None:
0413         run = FakeDictionary(
0414             {
0415                 "ALL_TAB_COLUMNS": [
0416                     ("PANDAID", "NUMBER"),
0417                     ("MODIFICATIONTIME", "DATE"),
0418                     ("JOBSTATUS", "VARCHAR2"),
0419                     ("STARTTIME", "TIMESTAMP(6)"),
0420                     ("JOBPARAMETERS", "CLOB"),
0421                     ("ROWID_X", "RAW"),
0422                 ],
0423                 # the hidden column of an index on an expression
0424                 "ALL_INDEXES": [("PANDAID", "LOCAL"), ("SYS_NC00034$", "LOCAL"), ("JOBSTATUS", None)],
0425                 "ALL_PART_TABLES": [("RANGE",)],
0426                 "ALL_PART_KEY_COLUMNS": [("MODIFICATIONTIME",)],
0427                 "ALL_TAB_PARTITIONS": [(oracle_date("2012-08-01"),), (oracle_date("2026-10-02"),), ("TIMESTAMP' 2026-10-03 00:00:00'",)],
0428             }
0429         )
0430         info = read_table_info(run, "oracle", "atlas_panda", "jobsArchived4")
0431         self.assertEqual((info.owner, info.table_name), ("ATLAS_PANDA", "JOBSARCHIVED4"))
0432         self.assertEqual(
0433             info.columns,
0434             {
0435                 "pandaid": KIND_NUMBER,
0436                 "modificationtime": KIND_TIME,
0437                 "jobstatus": KIND_TEXT,
0438                 "starttime": KIND_TIME,
0439                 "jobparameters": KIND_LOB,
0440                 "rowid_x": KIND_OTHER,
0441             },
0442         )
0443         self.assertEqual(info.indexes, (IndexInfo("pandaid", True), IndexInfo("jobstatus", False)))
0444         self.assertEqual(
0445             info.partitioning,
0446             PartitionInfo(
0447                 "modificationtime",
0448                 (
0449                     (None, datetime.datetime(2012, 8, 1)),
0450                     (datetime.datetime(2012, 8, 1), datetime.datetime(2026, 10, 2)),
0451                     (datetime.datetime(2026, 10, 2), datetime.datetime(2026, 10, 3)),
0452                 ),
0453             ),
0454         )
0455         self.assertEqual(run.statements[0][1], {":owner": "ATLAS_PANDA", ":table_name": "JOBSARCHIVED4"})
0456
0457     def test_oracle_number_partitions(self) -> None:
0458         run = FakeDictionary(
0459             {
0460                 "ALL_TAB_COLUMNS": [("JEDITASKID", "NUMBER"), ("STATUS", "VARCHAR2")],
0461                 "ALL_INDEXES": [("JEDITASKID", "LOCAL")],
0462                 "ALL_PART_TABLES": [("RANGE",)],
0463                 "ALL_PART_KEY_COLUMNS": [("JEDITASKID",)],
0464                 "ALL_TAB_PARTITIONS": [("500000",), ("1000000",), ("MAXVALUE",)],
0465             }
0466         )
0467         info = read_table_info(run, "oracle", "ATLAS_PANDA", "JEDI_Tasks")
0468         assert info.partitioning is not None
0469         self.assertEqual(info.partitioning.bounds, ((None, 500000), (500000, 1000000), (1000000, None)))
0470
0471     def test_oracle_other_partitioning(self) -> None:
0472         # list partitions, whose bounds say nothing of a range
0473         run = FakeDictionary(
0474             {
0475                 "ALL_TAB_COLUMNS": [("ID", "NUMBER"), ("VO", "VARCHAR2")],
0476                 "ALL_PART_TABLES": [("LIST",)],
0477                 "ALL_PART_KEY_COLUMNS": [("VO",)],
0478                 "ALL_TAB_PARTITIONS": [("'atlas'",), ("'other'",)],
0479             }
0480         )
0481         info = read_table_info(run, "oracle", "ATLAS_PANDA", "T")
0482         self.assertEqual(info.partitioning, PartitionInfo("vo", ((None, None), (None, None))))
0483
0484     def test_postgres(self) -> None:
0485         run = FakeDictionary(
0486             {
0487                 "information_schema.columns": [
0488                     ("pandaid", "bigint"),
0489                     ("modificationtime", "timestamp without time zone"),
0490                     ("jobstatus", "character varying"),
0491                     ("metadata", "jsonb"),
0492                 ],
0493                 "pg_class": [("p", "RANGE (modificationtime)")],
0494                 "pg_index": [("pandaid",)],
0495                 "pg_inherits": [("FOR VALUES FROM ('2026-10-01 00:00:00') TO ('2026-10-02 00:00:00')",), ("DEFAULT",)],
0496             }
0497         )
0498         info = read_table_info(run, "postgres", "DOMA_PANDA", "jobsArchived4")
0499         self.assertEqual((info.owner, info.table_name), ("doma_panda", "jobsarchived4"))
0500         self.assertEqual(info.columns["metadata"], KIND_LOB)
0501         self.assertEqual(info.indexes, (IndexInfo("pandaid", True),))
0502         self.assertEqual(
0503             info.partitioning,
0504             PartitionInfo("modificationtime", ((datetime.datetime(2026, 10, 1), datetime.datetime(2026, 10, 2)), (None, None))),
0505         )
0506         # what WrappedCursor makes of the statements: the two bind variables and nothing else
0507         self.assertEqual(len(run.statements), 4)
0508         for sql, binds in run.statements:
0509             converted, params = convert_query_in_printf_format(sql, [binds], {})
0510             self.assertNotIn(":", converted)
0511             self.assertEqual(params, [["doma_panda", "jobsarchived4"]])
0512
0513     def test_postgres_not_partitioned(self) -> None:
0514         run = FakeDictionary({"information_schema.columns": [("key", "character varying")], "pg_class": [("r", None)], "pg_index": [("key",)]})
0515         info = read_table_info(run, "postgres", "doma_panda", "config")
0516         self.assertEqual((info.indexes, info.partitioning), ((IndexInfo("key", False),), None))
0517
0518     def test_not_visible(self) -> None:
0519         with self.assertRaises(TableNotVisible):
0520             read_table_info(FakeDictionary({}), "oracle", "ATLAS_PANDAARCH", "jobsArchived")
0521
0522     def test_cache(self) -> None:
0523         run = FakeDictionary({"ALL_TAB_COLUMNS": [("ID", "NUMBER")]})
0524         with mock.patch.object(schema_info.time, "monotonic", return_value=1000.0):
0525             first = get_table_info(run, "oracle", "ATLAS_PANDA", "CONFIG")
0526             self.assertIs(get_table_info(run, "oracle", "atlas_panda", "config"), first)
0527         n_statements = len(run.statements)
0528         with mock.patch.object(schema_info.time, "monotonic", return_value=1000.0 + schema_info.CACHE_SECONDS):
0529             self.assertIsNot(get_table_info(run, "oracle", "ATLAS_PANDA", "CONFIG"), first)
0530         self.assertGreater(len(run.statements), n_statements)
0531
0532
0533 if __name__ == "__main__":
0534     unittest.main()