Coverage for src/pyadql/parser.py: 94%

1"""

2ADQL 2.1 parser -> typed AST.

4Usage:

5 from pyadql.parser import parse_adql

6 ast = parse_adql("SELECT TOP 10 ra, dec FROM mytable WHERE ra > 10")

7"""

9from __future__ import annotations

11import os

12import time

14from lark import Lark, Token, Transformer

15from loguru import logger

17from . import ast_nodes as A

19_GRAMMAR_DIR = os.path.join(os.path.dirname(__file__), "grammar")

20_GRAMMAR_ENTRY = os.path.join(_GRAMMAR_DIR, "adql.lark")

22_lark_parser: Lark | None = None # compiled lazily, see _get_lark_parser()

25def _get_lark_parser() -> Lark:

26 """Compile (once, on first call) and cache the Lark parser for the ADQL

27 grammar.

29 Deliberately lazy rather than run at module import time: compiling the

30 Lark grammar takes ~100 ms and emits useful DEBUG logs (duration,

31 grammar size) -- if that happened at `pyadql` import time (triggered by

32 a bare `import pyadql`, before an application has had a chance to

33 configure its loguru sinks), those logs would either be invisible or

34 emitted with loguru's default configuration, ignoring the verbosity

35 level requested by the caller (e.g. the CLI's -v/-vv/-q flags).

37 The grammar itself is split across several files under grammar/

38 (lexer.lark, literals.lark, core.lark, adql.lark) for readability --

39 see the header comment in grammar/core.lark for why the mutually

40 recursive "core" of the grammar couldn't be split further. Lark.open()

41 resolves the relative %import statements between those files against

42 grammar/adql.lark's own directory.

43 """

44 global _lark_parser

45 if _lark_parser is not None:

46 return _lark_parser

48 logger.debug("Loading the ADQL grammar from {}", _GRAMMAR_ENTRY)

50 # Earley: more tolerant of the ambiguities inherent to a SQL-like grammar

51 # than LALR, which matters for covering all of ADQL without having to

52 # manually resolve every ambiguity by hand.

53 t0 = time.perf_counter()

54 _lark_parser = Lark.open(

55 _GRAMMAR_ENTRY, rel_to=_GRAMMAR_ENTRY, parser="earley", ambiguity="resolve"

56 )

57 logger.debug(

58 "ADQL grammar compiled in {:.1f} ms (earley engine, multi-file grammar under {})",

59 (time.perf_counter() - t0) * 1000,

60 _GRAMMAR_DIR,

61 )

62 return _lark_parser

65# ---------------------------------------------------------------------------

66# Internal markers used while assembling composite rules (select_query,

67# select_expression, table_expression, joins, ...). Avoids introducing

68# public AST classes for pure construction details.

69# ---------------------------------------------------------------------------

70class _Marker:

71 def __init__(self, kind, value=None):

72 self.kind = kind

73 self.value = value

76# Marker kinds handled by select_query(), grouped by how they translate into

77# a Query field: DISTINCT/ALL set "distinct" to a fixed value, the others

78# copy their marker's value onto the same-named field.

79_SELECT_MARKER_CONSTANTS = {"DISTINCT": ("distinct", True), "ALL": ("distinct", False)}

80_SELECT_MARKER_VALUE_FIELDS = {

81 "TOP": "top",

82 "WHERE": "where",

83 "GROUP_BY": "group_by",

84 "HAVING": "having",

85}

88def _apply_select_marker(fields: dict, item: _Marker) -> None:

89 if item.kind in _SELECT_MARKER_CONSTANTS:

90 field, value = _SELECT_MARKER_CONSTANTS[item.kind]

91 fields[field] = value

92 elif item.kind in _SELECT_MARKER_VALUE_FIELDS:

93 fields[_SELECT_MARKER_VALUE_FIELDS[item.kind]] = item.value

96def _assign_select_or_from(fields: dict, item) -> None:

97 # `list` items are ambiguous between select_items and from_clause --

98 # from_clause always comes after select_list positionally, so the

99 # first list/Star we see is select_list and any subsequent list is

100 # from_clause.

101 if fields["select_list"] is None:

102 fields["select_list"] = item

103 else:

104 fields["from_clause"] = item

105

106

107def _unwrap_name(tok) -> str:

108 """Convert a NAME token into a plain Python identifier, handling

109 double-quoted (delimited, case-sensitive) identifiers, including the

110 doubled-double-quote escape for a literal " inside them

111 ([AnnexA #delimited_identifier], #double_quote_symbol)."""

112 s = str(tok)

113 if s.startswith('"') and s.endswith('"'):

114 return s[1:-1].replace('""', '"')

115 return s

116

117

118def _unwrap_string(tok) -> str:

119 s = str(tok)

120 # strip the surrounding SQL quotes and undo the doubled-quote escape ('')

121 return s[1:-1].replace("''", "'")

122

123

124class ADQLTransformer(Transformer):

125 # -- literals -------------------------------------------------------------

126 def number_literal(self, c):

127 return A.NumberLiteral(float(c[0]))

128

129 def string_literal(self, c):

130 return A.StringLiteral(_unwrap_string(c[0]))

131

132 def null_literal(self, c):

133 return A.NullLiteral()

134

135 def true_literal(self, c):

136 return A.BoolLiteral(True)

137

138 def false_literal(self, c):

139 return A.BoolLiteral(False)

140

141 # -- CAST target types ------------------------------------------------------

142 def string_cast_type(self, c):

143 name = str(c[0]).upper()

144 params = [int(t) for t in c[1:]] if len(c) > 1 else None

145 return A.CastType(name, params)

146

147 def numeric_cast_type(self, c):

148 return A.CastType(" ".join(str(t).upper() for t in c))

149

150 def datetime_cast_type(self, c):

151 return A.CastType(str(c[0]).upper())

152

153 def geometry_cast_type(self, c):

154 return A.CastType(str(c[0]).upper())

155

156 # -- identifiers ------------------------------------------------------------

157 def correlated_name(self, c):

158 return _unwrap_name(c[0])

159

160 def column_name(self, c):

161 return _unwrap_name(c[0])

162

163 def query_name(self, c):

164 return _unwrap_name(c[0])

165

166 def table_name(self, c):

167 return ".".join(_unwrap_name(t) for t in c)

168

169 def column_reference(self, c):

170 # lark.Token subclasses str: conversion must be forced for every

171 # segment, otherwise "raw" segments (t.<NAME>) stay Token instances

172 # in the repr, even though isinstance(x, str) is already true.

173 parts = [_unwrap_name(p) if isinstance(p, Token) else p for p in c]

174 return A.ColumnRef(parts)

175

176 def column_name_list(self, c):

177 return list(c)

178

179 # -- arithmetic / logical operators ---------------------------------------

180 def add(self, c):

181 return A.BinaryOp("+", c[0], c[1])

182

183 def sub(self, c):

184 return A.BinaryOp("-", c[0], c[1])

185

186 def mul(self, c):

187 return A.BinaryOp("*", c[0], c[1])

188

189 def div(self, c):

190 return A.BinaryOp("/", c[0], c[1])

191

192 def neg(self, c):

193 return A.UnaryOp("-", c[0])

194

195 def pos(self, c):

196 return A.UnaryOp("+", c[0])

197

198 def concat(self, c):

199 return A.BinaryOp("||", c[0], c[1])

200

201 def and_(self, c):

202 return A.BinaryOp("AND", c[0], c[1])

203

204 def or_(self, c):

205 return A.BinaryOp("OR", c[0], c[1])

206

207 def not_(self, c):

208 return A.UnaryOp("NOT", c[0])

209

210 def scalar_subquery(self, c):

211 return A.ScalarSubquery(c[0])

212

213 # -- predicates -------------------------------------------------------------

214 def comparison_predicate(self, c):

215 left, op, right = c[0], str(c[1]), c[2]

216 return A.BinaryOp(op, left, right)

217

218 def between_predicate(self, c):

219 negated = any(isinstance(t, Token) and t.type == "NOT_KW" for t in c)

220 vals = [x for x in c if not (isinstance(x, Token) and x.type == "NOT_KW")]

221 expr, low, high = vals

222 return A.Between(expr, low, high, negated)

223

224 def in_predicate(self, c):

225 negated = any(isinstance(t, Token) and t.type == "NOT_KW" for t in c)

226 vals = [x for x in c if not (isinstance(x, Token) and x.type == "NOT_KW")]

227 expr, values = vals[0], vals[1]

228 return A.InPredicate(expr, values, negated)

229

230 def in_predicate_value(self, c):

231 return c[0]

232

233 def in_subquery(self, c):

234 return c[0]

235

236 def in_value_list(self, c):

237 return list(c)

238

239 def like_predicate(self, c):

240 negated = any(isinstance(t, Token) and t.type == "NOT_KW" for t in c)

241 case_insensitive = any(isinstance(t, Token) and t.type == "ILIKE_KW" for t in c)

242 vals = [

243 x

244 for x in c

245 if not (

246 isinstance(x, Token) and x.type in ("NOT_KW", "LIKE_KW", "ILIKE_KW")

247 )

248 ]

249 expr, pattern = vals

250 return A.Like(expr, pattern, negated, case_insensitive)

251

252 def null_predicate(self, c):

253 negated = any(isinstance(t, Token) and t.type == "NOT_KW" for t in c)

254 expr = c[0]

255 return A.IsNull(expr, negated)

256

257 def exists_predicate(self, c):

258 return A.Exists(c[0])

259

260 # -- functions --------------------------------------------------------

261 def function_call(self, c):

262 return c[0]

263

264 def set_function(self, c):

265 name = str(c[0]).upper()

266 rest = c[1:]

267 distinct = False

268 args = []

269 for item in rest:

270 if isinstance(item, _Marker) and item.kind == "DISTINCT":

271 distinct = True

272 elif isinstance(item, _Marker) and item.kind == "ALL":

273 distinct = False

274 else:

275 args.append(item)

276 return A.FunctionCall(name, args, distinct)

277

278 def count_star(self, c):

279 return A.CountStar()

280

281 def numeric_function(self, c):

282 name = str(c[0]).upper()

283 args = list(c[1]) if len(c) > 1 else []

284 return A.FunctionCall(name, args, False)

285

286 def string_function(self, c):

287 name = str(c[0]).upper()

288 args = list(c[1]) if len(c) > 1 else []

289 return A.FunctionCall(name, args, False)

290

291 def udf(self, c):

292 name = c[0]

293 args = list(c[1]) if len(c) > 1 else []

294 return A.UserFunctionCall(name, args)

295

296 def arg_list(self, c):

297 return list(c)

298

299 def cast_specification(self, c):

300 return A.Cast(c[0], c[1])

301

302 def coalesce_expression(self, c):

303 return A.Coalesce(list(c))

304

305 # -- geometry ----------------------------------------------------------

306 def geometry_function(self, c):

307 return c[0]

308

309 def coord_sys(self, c):

310 return A.StringLiteral(_unwrap_string(c[0]))

311

312 def coordinates(self, c):

313 return A.Coordinates(c[0], c[1])

314

315 def radius(self, c):

316 return c[0]

317

318 def coord_value(self, c):

319 return c[0]

320

321 def udf_as_point(self, c):

322 name = c[0]

323 args = list(c[1]) if len(c) > 1 else []

324 return A.UserFunctionCall(name, args)

325

326 def point_expr(self, c):

327 # children: [coord_sys?] coordinates

328 if len(c) == 2:

329 coordsys, coords = c

330 else:

331 coordsys, coords = None, c[0]

332 return A.Point(coordsys, coords.ra, coords.dec)

333

334 def circle_center(self, c):

335 return c[0]

336

337 def circle_expr(self, c):

338 # children: [coord_sys?] circle_center radius

339 if len(c) == 3:

340 coordsys, center, radius = c

341 else:

342 coordsys = None

343 center, radius = c

344 return A.Circle(coordsys, center, radius)

345

346 def box_center(self, c):

347 return c[0]

348

349 def box_expr(self, c):

350 # children: [coord_sys?] box_center width height

351 if len(c) == 4:

352 coordsys, center, width, height = c

353 else:

354 coordsys = None

355 center, width, height = c

356 return A.Box(coordsys, center, width, height)

357

358 def polygon_vertices_coords(self, c):

359 return list(c)

360

361 def polygon_vertices_points(self, c):

362 return list(c)

363

364 def polygon_expr(self, c):

365 # children: [coord_sys?] polygon_vertices(list)

366 if len(c) == 2:

367 coordsys, vertices = c

368 else:

369 coordsys, vertices = None, c[0]

370 return A.Polygon(coordsys, vertices)

371

372 def region_expr(self, c):

373 return A.Region(_unwrap_string(c[0]))

374

375 def centroid_expr(self, c):

376 return A.Centroid(c[0])

377

378 def area_expr(self, c):

379 return A.Area(c[0])

380

381 def coord1_expr(self, c):

382 return A.Coord1(c[0])

383

384 def coord2_expr(self, c):

385 return A.Coord2(c[0])

386

387 def coordsys_expr(self, c):

388 return A.Coordsys(c[0])

389

390 def distance_points(self, c):

391 return A.Distance(list(c), numeric_form=False)

392

393 def distance_coords(self, c):

394 return A.Distance(list(c), numeric_form=True)

395

396 def contains_expr(self, c):

397 return A.Contains(c[0], c[1])

398

399 def intersects_expr(self, c):

400 return A.Intersects(c[0], c[1])

401

402 # -- SELECT / clauses -----------------------------------------------------

403 def distinct(self, c):

404 return _Marker("DISTINCT")

405

406 def all_(self, c):

407 return _Marker("ALL")

408

409 def top(self, c):

410 return _Marker("TOP", int(c[0]))

411

412 def star_select(self, c):

413 return A.Star()

414

415 def qualified_star(self, c):

416 return A.Star() # qualified star (t.*) simplified to a generic Star

417

418 def alias(self, c):

419 return _Marker("ALIAS", c[0])

420

421 def derived_column(self, c):

422 expr = c[0]

423 alias = None

424 if len(c) > 1 and isinstance(c[1], _Marker) and c[1].kind == "ALIAS":

425 alias = c[1].value

426 return A.SelectItem(expr, alias)

427

428 def select_sublist(self, c):

429 return c[0]

430

431 def select_items(self, c):

432 return list(c)

433

434 def asc(self, c):

435 return "ASC"

436

437 def desc(self, c):

438 return "DESC"

439

440 def sort_spec(self, c):

441 expr = c[0]

442 direction = c[1] if len(c) > 1 else "ASC"

443 return A.SortItem(expr, direction)

444

445 def where_clause(self, c):

446 return _Marker("WHERE", c[0])

447

448 def group_by_clause(self, c):

449 return _Marker("GROUP_BY", list(c))

450

451 def having_clause(self, c):

452 return _Marker("HAVING", c[0])

453

454 def order_by_clause(self, c):

455 return _Marker("ORDER_BY", list(c))

456

457 def offset_clause(self, c):

458 return _Marker("OFFSET", int(c[0]))

459

460 def from_clause(self, c):

461 return list(c)

462

463 def select_query(self, c):

464 fields = {

465 "distinct": False,

466 "top": None,

467 "select_list": None,

468 "from_clause": [],

469 "where": None,

470 "group_by": None,

471 "having": None,

472 }

473 for item in c:

474 if isinstance(item, _Marker):

475 _apply_select_marker(fields, item)

476 elif isinstance(item, (A.Star, list)):

477 _assign_select_or_from(fields, item)

478 return A.Query(**fields)

479

480 # -- set operations / select_expression -----------------------------------

481 def _set_op(self, op_name, c):

482 left = c[0]

483 right = c[-1]

484 distinct = not any(

485 isinstance(item, Token) and item.type == "ALL_KW" for item in c[1:-1]

486 )

487 return A.SetOperation(op_name, distinct, left, right)

488

489 def union_op(self, c):

490 return self._set_op("UNION", c)

491

492 def except_op(self, c):

493 return self._set_op("EXCEPT", c)

494

495 def intersect_op(self, c):

496 return self._set_op("INTERSECT", c)

497

498 def paren_select_expression(self, c):

499 return c[0]

500

501 def select_expression(self, c):

502 body = c[0]

503 order_by = None

504 offset = None

505 for item in c[1:]:

506 if isinstance(item, _Marker) and item.kind == "ORDER_BY":

507 order_by = item.value

508 elif isinstance(item, _Marker) and item.kind == "OFFSET":

509 offset = item.value

510 return A.SelectExpression(body=body, order_by=order_by, offset=offset)

511

512 def with_query(self, c):

513 return A.CTE(c[0], c[1])

514

515 def with_clause(self, c):

516 return _Marker("WITH", list(c))

517

518 def query_specification(self, c):

519 with_clause = None

520 select_expression = None

521 for item in c:

522 if isinstance(item, _Marker) and item.kind == "WITH":

523 with_clause = item.value

524 elif isinstance(item, A.SelectExpression):

525 select_expression = item

526 select_expression.with_clause = with_clause

527 return select_expression

528

529 # -- FROM / joins -----------------------------------------------------

530 def table_ref(self, c):

531 name = c[0]

532 alias, cols = None, None

533 if len(c) > 1 and c[1] is not None:

534 alias, cols = c[1]

535 return A.TableRef(name, alias, cols)

536

537 def derived_table(self, c):

538 subquery = c[0]

539 alias, cols = None, None

540 if len(c) > 1 and c[1] is not None:

541 alias, cols = c[1]

542 return A.DerivedTable(subquery, alias, cols)

543

544 def paren_table_ref(self, c):

545 return c[0]

546

547 def correlation_specification(self, c):

548 alias = c[0]

549 cols = c[1] if len(c) > 1 else None

550 return (alias, cols)

551

552 def inner_join(self, c):

553 return "INNER"

554

555 def left_join(self, c):

556 return "LEFT"

557

558 def right_join(self, c):

559 return "RIGHT"

560

561 def full_join(self, c):

562 return "FULL"

563

564 def on_spec(self, c):

565 return _Marker("ON", c[0])

566

567 def using_spec(self, c):

568 return _Marker("USING", list(c))

569

570 def joined(self, c):

571 natural = False

572 join_type = "INNER"

573 right = None

574 on = None

575 using = None

576 for item in c:

577 if isinstance(item, Token) and item.upper() == "NATURAL":

578 natural = True

579 elif item in ("INNER", "LEFT", "RIGHT", "FULL"):

580 join_type = item

581 elif isinstance(item, _Marker) and item.kind == "ON":

582 on = item.value

583 elif isinstance(item, _Marker) and item.kind == "USING":

584 using = item.value

585 elif isinstance(item, (A.TableRef, A.DerivedTable, A.Join)):

586 right = item

587 return _Marker("JOINED", (natural, join_type, right, on, using))

588

589 def table_reference(self, c):

590 base = c[0]

591 for jm in c[1:]:

592 natural, join_type, right, on, using = jm.value

593 base = A.Join(base, right, join_type, natural, on, using)

594 return base

595

596 def start(self, c):

597 return c[0]

598

599

600_transformer = ADQLTransformer()

601

602

603def _summarize_ast(ast) -> str:

604 """Compact summary of an AST for DEBUG logs (avoids dumping the whole

605 tree when the goal is just to confirm its general shape)."""

606 if not isinstance(ast, A.SelectExpression):

607 return type(ast).__name__

608 body = ast.body

609 if isinstance(body, A.SetOperation):

610 body_desc = f"SetOperation(op={body.op}, distinct={body.distinct})"

611 elif isinstance(body, A.Query):

612 n_select = 1 if isinstance(body.select_list, A.Star) else len(body.select_list)

613 body_desc = (

614 f"Query(distinct={body.distinct}, top={body.top}, "

615 f"n_select={n_select}, n_from={len(body.from_clause)}, "

616 f"where={body.where is not None}, group_by={body.group_by is not None}, "

617 f"having={body.having is not None})"

618 )

619 else:

620 body_desc = type(body).__name__

621 return (

622 f"SelectExpression(with_clause={ast.with_clause is not None}, "

623 f"order_by={ast.order_by is not None}, offset={ast.offset}, body={body_desc})"

624 )

625

626

627def parse_adql(query: str):

628 """Parse an ADQL query and return the root of the typed AST

629 (an ast_nodes.SelectExpression; see that class's docstring for the full

630 shape)."""

631 logger.debug("parse_adql: query of {} characters", len(query))

632

633 t0 = time.perf_counter()

634 try:

635 tree = _get_lark_parser().parse(query)

636 except Exception:

637 logger.debug(

638 "Failed at the grammar stage (Lark) after {:.1f} ms -- the query "

639 "does not follow the expected ADQL/SQL syntax",

640 (time.perf_counter() - t0) * 1000,

641 )

642 raise

643 t_parse = time.perf_counter()

644 logger.debug(

645 "Lark tree obtained in {:.1f} ms ({} direct children under 'start')",

646 (t_parse - t0) * 1000,

647 len(tree.children),

648 )

649 logger.opt(lazy=True).trace("Raw Lark tree:\n{}", lambda: tree.pretty())

650

651 try:

652 ast = _transformer.transform(tree)

653 except Exception:

654 logger.debug(

655 "Failed at the transformer stage (Lark tree -> typed AST) after "

656 "{:.1f} ms -- the grammar matched but building the AST node "

657 "failed (check ADQLTransformer in parser.py)",

658 (time.perf_counter() - t_parse) * 1000,

659 )

660 raise

661 logger.debug(

662 "AST built in {:.1f} ms (total parse+transform: {:.1f} ms): {}",

663 (time.perf_counter() - t_parse) * 1000,

664 (time.perf_counter() - t0) * 1000,

665 _summarize_ast(ast),

666 )

667 return ast

668

669

670def parse_tree(query: str):

671 """Return the raw Lark parse tree (useful for debugging the grammar)."""

672 logger.debug("parse_tree: query of {} characters", len(query))

673 t0 = time.perf_counter()

674 tree = _get_lark_parser().parse(query)

675 logger.debug("Lark tree obtained in {:.1f} ms", (time.perf_counter() - t0) * 1000)

676 return tree