asyncprocess/db_read/schema_info.py

0001 """
0002 What the async db_read requests need to know of a table, as the data dictionary tells it.
0003 """
0004
0005 import datetime
0006 import re
0007 import threading
0008 import time
0009 from collections.abc import Callable
0010 from dataclasses import dataclass
0011 from typing import Any
0012
0013 from pandaserver.taskbuffer.db_proxy_mods.base_module import parse_partition_bounds
0014
0015 # a partition bound: a time for a table partitioned on a time column, an int for one
0016 # partitioned on a number, such as the JEDI tables on jediTaskID
0017 Bound = datetime.datetime | int
0018
0019 # kinds of column. The kind decides how a value in a request is read, and what a request
0020 # may do with the column
0021 KIND_NUMBER = "number"
0022 KIND_TEXT = "text"
0023 KIND_TIME = "time"
0024 # CLOB and the like: returned only when named, and never compared
0025 KIND_LOB = "lob"
0026 # anything else, such as RAW: returned, and never compared
0027 KIND_OTHER = "other"
0028
0029
0030 @dataclass(frozen=True)
0031 class IndexInfo:
0032     # lower-case name of the first column. An index led by an expression is not kept, since no
0033     # condition of a request can use it
0034     leading_column: str
0035     # True for an index partitioned like its table, which a lookup probes once per partition
0036     # touched; False for a global index or one on a table that is not partitioned
0037     local: bool
0038
0039
0040 @dataclass(frozen=True)
0041 class PartitionInfo:
0042     # lower-case name of the partition key column
0043     key_column: str
0044     # [lower, upper) of each partition, None where it is open
0045     bounds: tuple[tuple[Bound | None, Bound | None], ...]
0046
0047
0048 @dataclass(frozen=True)
0049 class TableInfo:
0050     owner: str
0051     table_name: str
0052     # lower-case column names in the order of the table, each with its kind
0053     columns: dict[str, str]
0054     indexes: tuple[IndexInfo, ...]
0055     # None when the table is not partitioned
0056     partitioning: PartitionInfo | None
0057
0058
0059 class TableNotVisible(Exception):
0060     """The data dictionary shows no columns of the table, which is what a missing grant looks like."""
0061
0062
0063 # runs one statement of the data dictionary and returns its rows
0064 RunStatement = Callable[[str, dict[str, Any]], list[tuple[Any, ...]]]
0065
0066 # how long a table read from the dictionary is kept, which is how long a new index or partition
0067 # may go unseen
0068 CACHE_SECONDS = 600
0069
0070 # the comment of every statement of the async db_read requests, their reads of the dictionary
0071 # included, which tells them apart in the database's views of what runs. It has neither ":" nor
0072 # "%", which the Postgres wrapper would take for binds
0073 STATEMENT_COMMENT = " /* async_process.db_read */"
0074
0075 _ORACLE_KINDS = {
0076     "NUMBER": KIND_NUMBER,
0077     "FLOAT": KIND_NUMBER,
0078     "INTEGER": KIND_NUMBER,
0079     "BINARY_FLOAT": KIND_NUMBER,
0080     "BINARY_DOUBLE": KIND_NUMBER,
0081     "VARCHAR2": KIND_TEXT,
0082     "NVARCHAR2": KIND_TEXT,
0083     "CHAR": KIND_TEXT,
0084     "NCHAR": KIND_TEXT,
0085     "DATE": KIND_TIME,
0086     "CLOB": KIND_LOB,
0087     "NCLOB": KIND_LOB,
0088     "BLOB": KIND_LOB,
0089     "LONG": KIND_LOB,
0090 }
0091
0092 _POSTGRES_KINDS = {
0093     "smallint": KIND_NUMBER,
0094     "integer": KIND_NUMBER,
0095     "bigint": KIND_NUMBER,
0096     "numeric": KIND_NUMBER,
0097     "real": KIND_NUMBER,
0098     "double precision": KIND_NUMBER,
0099     "character varying": KIND_TEXT,
0100     "character": KIND_TEXT,
0101     "text": KIND_TEXT,
0102     "date": KIND_TIME,
0103     "timestamp without time zone": KIND_TIME,
0104     "timestamp with time zone": KIND_TIME,
0105     "json": KIND_LOB,
0106     "jsonb": KIND_LOB,
0107 }
0108
0109 _INTEGER = re.compile(r"^\s*(-?\d+)\s*$")
0110 # FOR VALUES FROM (lower) TO (upper), as pg_get_expr spells a range partition
0111 _POSTGRES_RANGE = re.compile(r"^FOR VALUES FROM \((.*)\) TO \((.*)\)$")
0112 # RANGE (column), as pg_get_partkeydef spells a range partitioning on one column
0113 _POSTGRES_RANGE_KEY = re.compile(r"^RANGE \(([A-Za-z_][A-Za-z0-9_]*)\)$")
0114
0115 _cache: dict[tuple[str, str, str], tuple[float, TableInfo]] = {}
0116 _cache_lock = threading.Lock()
0117
0118
0119 def column_kind(backend: str, data_type: str) -> str:
0120     """
0121     Get the kind of a column from its data type as the dictionary spells it.
0122
0123     Args:
0124         backend(str): oracle or postgres
0125         data_type(str): the data type, such as TIMESTAMP(6) for Oracle or bigint for Postgres
0126
0127     Returns:
0128         str: one of the KIND_ constants
0129     """
0130     if backend == "postgres":
0131         return _POSTGRES_KINDS.get(data_type, KIND_OTHER)
0132     if data_type.startswith("TIMESTAMP"):
0133         return KIND_TIME
0134     return _ORACLE_KINDS.get(data_type, KIND_OTHER)
0135
0136
0137 def parse_bound(text: str | None, kind: str) -> Bound | None:
0138     """
0139     Read one partition bound.
0140
0141     Args:
0142         text(str): the bound, an Oracle HIGH_VALUE or one side of a Postgres FOR VALUES
0143         kind(str): kind of the partition key column
0144
0145     Returns:
0146         datetime or int or None: the bound, or None for MAXVALUE, MINVALUE and what cannot be read,
0147             which leaves that side of the partition open
0148     """
0149     if kind == KIND_TIME:
0150         stamps = parse_partition_bounds(text)
0151         return stamps[0] if len(stamps) == 1 else None
0152     if kind == KIND_NUMBER:
0153         match = _INTEGER.match((text or "").strip("'"))
0154         return int(match.group(1)) if match else None
0155     return None
0156
0157
0158 def read_table_info(run: RunStatement, backend: str, owner: str, table_name: str) -> TableInfo:
0159     """
0160     Read a table from the data dictionary.
0161
0162     Args:
0163         run(callable): runs a statement and returns its rows; the caller puts each round trip under its limits
0164         backend(str): oracle or postgres
0165         owner(str): schema owning the table
0166         table_name(str): name of the table
0167
0168     Returns:
0169         TableInfo: the table
0170
0171     Raises:
0172         TableNotVisible: when the dictionary shows no columns of the table
0173     """
0174     if backend == "postgres":
0175         return _read_postgres(run, owner.lower(), table_name.lower())
0176     return _read_oracle(run, owner.upper(), table_name.upper())
0177
0178
0179 def _read_oracle(run: RunStatement, owner: str, table_name: str) -> TableInfo:
0180     names = {":owner": owner, ":table_name": table_name}
0181     rows = run("SELECT column_name, data_type FROM ALL_TAB_COLUMNS WHERE owner=:owner AND table_name=:table_name ORDER BY column_id" + STATEMENT_COMMENT, names)
0182     if not rows:
0183         raise TableNotVisible(f"{owner}.{table_name}")
0184     columns = {name.lower(): column_kind("oracle", data_type) for name, data_type in rows}
0185     # an index led by an expression is led by a hidden column such as SYS_NC00034$, which is not
0186     # among the columns above, and is left out with them
0187     rows = run(
0188         "SELECT c.column_name, p.locality FROM ALL_INDEXES i "
0189         "JOIN ALL_IND_COLUMNS c ON c.index_owner=i.owner AND c.index_name=i.index_name AND c.column_position=1 "
0190         "LEFT JOIN ALL_PART_INDEXES p ON p.owner=i.owner AND p.index_name=i.index_name "
0191         "WHERE i.table_owner=:owner AND i.table_name=:table_name " + STATEMENT_COMMENT,
0192         names,
0193     )
0194     indexes = tuple(IndexInfo(column.lower(), locality == "LOCAL") for column, locality in rows if column.lower() in columns)
0195     partitioning = None
0196     rows = run("SELECT partitioning_type FROM ALL_PART_TABLES WHERE owner=:owner AND table_name=:table_name " + STATEMENT_COMMENT, names)
0197     if rows:
0198         partitioning_type = rows[0][0]
0199         key_rows = run(
0200             "SELECT column_name FROM ALL_PART_KEY_COLUMNS WHERE owner=:owner AND name=:table_name AND object_type='TABLE' ORDER BY column_position "
0201             + STATEMENT_COMMENT,
0202             names,
0203         )
0204         bound_rows = run(
0205             "SELECT high_value FROM ALL_TAB_PARTITIONS WHERE table_owner=:owner AND table_name=:table_name ORDER BY partition_position " + STATEMENT_COMMENT,
0206             names,
0207         )
0208         key_column = key_rows[0][0].lower() if key_rows else ""
0209         # bounds are read for a range on one column, interval partitioning being a kind of range;
0210         # for anything else every partition is taken to hold any value
0211         readable = partitioning_type == "RANGE" and len(key_rows) == 1 and key_column in columns
0212         bounds: list[tuple[Bound | None, Bound | None]] = []
0213         lower: Bound | None = None
0214         for (high_value,) in bound_rows:
0215             if not readable:
0216                 bounds.append((None, None))
0217                 continue
0218             # a partition starts where the one before it ends
0219             upper = parse_bound(high_value, columns[key_column])
0220             bounds.append((lower, upper))
0221             lower = upper
0222         partitioning = PartitionInfo(key_column, tuple(bounds))
0223     return TableInfo(owner, table_name, columns, indexes, partitioning)
0224
0225
0226 def _read_postgres(run: RunStatement, owner: str, table_name: str) -> TableInfo:
0227     names = {":owner": owner, ":table_name": table_name}
0228     rows = run(
0229         "SELECT column_name, data_type FROM information_schema.columns WHERE table_schema=:owner AND table_name=:table_name ORDER BY ordinal_position "
0230         + STATEMENT_COMMENT,
0231         names,
0232     )
0233     if not rows:
0234         raise TableNotVisible(f"{owner}.{table_name}")
0235     columns = {name.lower(): column_kind("postgres", data_type) for name, data_type in rows}
0236     rows = run(
0237         "SELECT t.relkind, pg_get_partkeydef(t.oid) FROM pg_class t JOIN pg_namespace n ON n.oid=t.relnamespace "
0238         "WHERE n.nspname=:owner AND t.relname=:table_name " + STATEMENT_COMMENT,
0239         names,
0240     )
0241     relkind, partition_key = rows[0] if rows else ("r", None)
0242     # an index on a partitioned table is partitioned with it: Postgres has no other kind. An
0243     # index led by an expression has 0 as its first column, which no column matches
0244     rows = run(
0245         "SELECT a.attname FROM pg_index x JOIN pg_class t ON t.oid=x.indrelid JOIN pg_namespace n ON n.oid=t.relnamespace "
0246         "JOIN pg_attribute a ON a.attrelid=t.oid AND a.attnum=x.indkey[0] "
0247         "WHERE n.nspname=:owner AND t.relname=:table_name " + STATEMENT_COMMENT,
0248         names,
0249     )
0250     indexes = tuple(IndexInfo(column.lower(), relkind == "p") for (column,) in rows if column.lower() in columns)
0251     partitioning = None
0252     if relkind == "p":
0253         match = _POSTGRES_RANGE_KEY.match(partition_key or "")
0254         key_column = match.group(1).lower() if match else ""
0255         readable = key_column in columns
0256         rows = run(
0257             "SELECT pg_get_expr(c.relpartbound, c.oid) FROM pg_inherits i "
0258             "JOIN pg_class c ON c.oid=i.inhrelid JOIN pg_class p ON p.oid=i.inhparent "
0259             "JOIN pg_namespace n ON n.oid=p.relnamespace "
0260             "WHERE n.nspname=:owner AND p.relname=:table_name " + STATEMENT_COMMENT,
0261             names,
0262         )
0263         bounds: list[tuple[Bound | None, Bound | None]] = []
0264         for (bound,) in rows:
0265             # the default partition, and the bound of anything but a range, take any value
0266             sides = _POSTGRES_RANGE.match(bound or "") if readable else None
0267             if sides is None:
0268                 bounds.append((None, None))
0269             else:
0270                 bounds.append((parse_bound(sides.group(1), columns[key_column]), parse_bound(sides.group(2), columns[key_column])))
0271         partitioning = PartitionInfo(key_column, tuple(bounds))
0272     return TableInfo(owner, table_name, columns, indexes, partitioning)
0273
0274
0275 def get_table_info(run: RunStatement, backend: str, owner: str, table_name: str) -> TableInfo:
0276     """
0277     Get a table from the data dictionary, read again after CACHE_SECONDS.
0278
0279     Args:
0280         run(callable): runs a statement and returns its rows, used when the table is read
0281         backend(str): oracle or postgres
0282         owner(str): schema owning the table
0283         table_name(str): name of the table
0284
0285     Returns:
0286         TableInfo: the table
0287
0288     Raises:
0289         TableNotVisible: when the dictionary shows no columns of the table
0290     """
0291     key = (backend, owner.lower(), table_name.lower())
0292     with _cache_lock:
0293         cached = _cache.get(key)
0294     if cached is not None and time.monotonic() - cached[0] < CACHE_SECONDS:
0295         return cached[1]
0296     info = read_table_info(run, backend, owner, table_name)
0297     with _cache_lock:
0298         _cache[key] = (time.monotonic(), info)
0299     return info