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

76def _unwrap_name(tok) -> str:

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

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

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

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

81 s = str(tok)

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

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

84 return s

87def _unwrap_string(tok) -> str:

88 s = str(tok)

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

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

93class ADQLTransformer(Transformer):

94 # -- literals -------------------------------------------------------------

95 def number_literal(self, c):

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

98 def string_literal(self, c):

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

100

101 def null_literal(self, c):

102 return A.NullLiteral()

103

104 def true_literal(self, c):

105 return A.BoolLiteral(True)

106

107 def false_literal(self, c):

108 return A.BoolLiteral(False)

109

110 # -- CAST target types ------------------------------------------------------

111 def string_cast_type(self, c):

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

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

114 return A.CastType(name, params)

115

116 def numeric_cast_type(self, c):

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

118

119 def datetime_cast_type(self, c):

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

121

122 def geometry_cast_type(self, c):

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

124

125 # -- identifiers ------------------------------------------------------------

126 def correlated_name(self, c):

127 return _unwrap_name(c[0])

128

129 def column_name(self, c):

130 return _unwrap_name(c[0])

131

132 def query_name(self, c):

133 return _unwrap_name(c[0])

134

135 def table_name(self, c):

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

137

138 def column_reference(self, c):

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

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

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

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

143 return A.ColumnRef(parts)

144

145 def column_name_list(self, c):

146 return list(c)

147

148 # -- arithmetic / logical operators ---------------------------------------

149 def add(self, c):

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

151

152 def sub(self, c):

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

154

155 def mul(self, c):

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

157

158 def div(self, c):

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

160

161 def neg(self, c):

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

163

164 def pos(self, c):

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

166

167 def concat(self, c):

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

169

170 def and_(self, c):

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

172

173 def or_(self, c):

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

175

176 def not_(self, c):

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

178

179 def scalar_subquery(self, c):

180 return A.ScalarSubquery(c[0])

181

182 # -- predicates -------------------------------------------------------------

183 def comparison_predicate(self, c):

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

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

186

187 def between_predicate(self, c):

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

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

190 expr, low, high = vals

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

192

193 def in_predicate(self, c):

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

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

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

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

198

199 def in_predicate_value(self, c):

200 return c[0]

201

202 def in_subquery(self, c):

203 return c[0]

204

205 def in_value_list(self, c):

206 return list(c)

207

208 def like_predicate(self, c):

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

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

211 vals = [

212 x

213 for x in c

214 if not (

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

216 )

217 ]

218 expr, pattern = vals

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

220

221 def null_predicate(self, c):

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

223 expr = c[0]

224 return A.IsNull(expr, negated)

225

226 def exists_predicate(self, c):

227 return A.Exists(c[0])

228

229 # -- functions --------------------------------------------------------

230 def function_call(self, c):

231 return c[0]

232

233 def set_function(self, c):

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

235 rest = c[1:]

236 distinct = False

237 args = []

238 for item in rest:

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

240 distinct = True

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

242 distinct = False

243 else:

244 args.append(item)

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

246

247 def count_star(self, c):

248 return A.CountStar()

249

250 def numeric_function(self, c):

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

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

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

254

255 def string_function(self, c):

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

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

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

259

260 def udf(self, c):

261 name = c[0]

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

263 return A.UserFunctionCall(name, args)

264

265 def arg_list(self, c):

266 return list(c)

267

268 def cast_specification(self, c):

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

270

271 def coalesce_expression(self, c):

272 return A.Coalesce(list(c))

273

274 # -- geometry ----------------------------------------------------------

275 def geometry_function(self, c):

276 return c[0]

277

278 def coord_sys(self, c):

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

280

281 def coordinates(self, c):

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

283

284 def radius(self, c):

285 return c[0]

286

287 def coord_value(self, c):

288 return c[0]

289

290 def udf_as_point(self, c):

291 name = c[0]

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

293 return A.UserFunctionCall(name, args)

294

295 def point_expr(self, c):

296 # children: [coord_sys?] coordinates

297 if len(c) == 2:

298 coordsys, coords = c

299 else:

300 coordsys, coords = None, c[0]

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

302

303 def circle_center(self, c):

304 return c[0]

305

306 def circle_expr(self, c):

307 # children: [coord_sys?] circle_center radius

308 if len(c) == 3:

309 coordsys, center, radius = c

310 else:

311 coordsys = None

312 center, radius = c

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

314

315 def box_center(self, c):

316 return c[0]

317

318 def box_expr(self, c):

319 # children: [coord_sys?] box_center width height

320 if len(c) == 4:

321 coordsys, center, width, height = c

322 else:

323 coordsys = None

324 center, width, height = c

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

326

327 def polygon_vertices_coords(self, c):

328 return list(c)

329

330 def polygon_vertices_points(self, c):

331 return list(c)

332

333 def polygon_expr(self, c):

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

335 if len(c) == 2:

336 coordsys, vertices = c

337 else:

338 coordsys, vertices = None, c[0]

339 return A.Polygon(coordsys, vertices)

340

341 def region_expr(self, c):

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

343

344 def centroid_expr(self, c):

345 return A.Centroid(c[0])

346

347 def area_expr(self, c):

348 return A.Area(c[0])

349

350 def coord1_expr(self, c):

351 return A.Coord1(c[0])

352

353 def coord2_expr(self, c):

354 return A.Coord2(c[0])

355

356 def coordsys_expr(self, c):

357 return A.Coordsys(c[0])

358

359 def distance_points(self, c):

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

361

362 def distance_coords(self, c):

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

364

365 def contains_expr(self, c):

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

367

368 def intersects_expr(self, c):

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

370

371 # -- SELECT / clauses -----------------------------------------------------

372 def distinct(self, c):

373 return _Marker("DISTINCT")

374

375 def all_(self, c):

376 return _Marker("ALL")

377

378 def top(self, c):

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

380

381 def star_select(self, c):

382 return A.Star()

383

384 def qualified_star(self, c):

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

386

387 def alias(self, c):

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

389

390 def derived_column(self, c):

391 expr = c[0]

392 alias = None

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

394 alias = c[1].value

395 return A.SelectItem(expr, alias)

396

397 def select_sublist(self, c):

398 return c[0]

399

400 def select_items(self, c):

401 return list(c)

402

403 def asc(self, c):

404 return "ASC"

405

406 def desc(self, c):

407 return "DESC"

408

409 def sort_spec(self, c):

410 expr = c[0]

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

412 return A.SortItem(expr, direction)

413

414 def where_clause(self, c):

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

416

417 def group_by_clause(self, c):

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

419

420 def having_clause(self, c):

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

422

423 def order_by_clause(self, c):

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

425

426 def offset_clause(self, c):

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

428

429 def from_clause(self, c):

430 return list(c)

431

432 def select_query(self, c):

433 distinct = False

434 top = None

435 select_list = None

436 from_clause = []

437 where = None

438 group_by = None

439 having = None

440 for item in c:

441 if isinstance(item, _Marker):

442 if item.kind == "DISTINCT":

443 distinct = True

444 elif item.kind == "ALL":

445 distinct = False

446 elif item.kind == "TOP":

447 top = item.value

448 elif item.kind == "WHERE":

449 where = item.value

450 elif item.kind == "GROUP_BY":

451 group_by = item.value

452 elif item.kind == "HAVING":

453 having = item.value

454 elif isinstance(item, A.Star) or isinstance(item, list):

455 # `list` items are ambiguous between select_items and

456 # from_clause -- from_clause always comes after select_list

457 # positionally, so the first list/Star we see is select_list

458 # and any subsequent list is from_clause.

459 if select_list is None:

460 select_list = item

461 else:

462 from_clause = item

463 return A.Query(

464 distinct=distinct,

465 top=top,

466 select_list=select_list,

467 from_clause=from_clause,

468 where=where,

469 group_by=group_by,

470 having=having,

471 )

472

473 # -- set operations / select_expression -----------------------------------

474 def _set_op(self, op_name, c):

475 left = c[0]

476 right = c[-1]

477 distinct = not any(

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

479 )

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

481

482 def union_op(self, c):

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

484

485 def except_op(self, c):

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

487

488 def intersect_op(self, c):

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

490

491 def paren_select_expression(self, c):

492 return c[0]

493

494 def select_expression(self, c):

495 body = c[0]

496 order_by = None

497 offset = None

498 for item in c[1:]:

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

500 order_by = item.value

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

502 offset = item.value

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

504

505 def with_query(self, c):

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

507

508 def with_clause(self, c):

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

510

511 def query_specification(self, c):

512 with_clause = None

513 select_expression = None

514 for item in c:

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

516 with_clause = item.value

517 elif isinstance(item, A.SelectExpression):

518 select_expression = item

519 select_expression.with_clause = with_clause

520 return select_expression

521

522 # -- FROM / joins -----------------------------------------------------

523 def table_ref(self, c):

524 name = c[0]

525 alias, cols = None, None

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

527 alias, cols = c[1]

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

529

530 def derived_table(self, c):

531 subquery = 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.DerivedTable(subquery, alias, cols)

536

537 def paren_table_ref(self, c):

538 return c[0]

539

540 def correlation_specification(self, c):

541 alias = c[0]

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

543 return (alias, cols)

544

545 def inner_join(self, c):

546 return "INNER"

547

548 def left_join(self, c):

549 return "LEFT"

550

551 def right_join(self, c):

552 return "RIGHT"

553

554 def full_join(self, c):

555 return "FULL"

556

557 def on_spec(self, c):

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

559

560 def using_spec(self, c):

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

562

563 def joined(self, c):

564 natural = False

565 join_type = "INNER"

566 right = None

567 on = None

568 using = None

569 for item in c:

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

571 natural = True

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

573 join_type = item

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

575 on = item.value

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

577 using = item.value

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

579 right = item

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

581

582 def table_reference(self, c):

583 base = c[0]

584 for jm in c[1:]:

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

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

587 return base

588

589 def start(self, c):

590 return c[0]

591

592

593_transformer = ADQLTransformer()

594

595

596def _summarize_ast(ast) -> str:

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

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

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

600 return type(ast).__name__

601 body = ast.body

602 if isinstance(body, A.SetOperation):

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

604 elif isinstance(body, A.Query):

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

606 body_desc = (

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

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

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

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

611 )

612 else:

613 body_desc = type(body).__name__

614 return (

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

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

617 )

618

619

620def parse_adql(query: str):

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

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

623 shape)."""

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

625

626 t0 = time.perf_counter()

627 try:

628 tree = _get_lark_parser().parse(query)

629 except Exception:

630 logger.debug(

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

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

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

634 )

635 raise

636 t_parse = time.perf_counter()

637 logger.debug(

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

639 (t_parse - t0) * 1000,

640 len(tree.children),

641 )

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

643

644 try:

645 ast = _transformer.transform(tree)

646 except Exception:

647 logger.debug(

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

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

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

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

652 )

653 raise

654 logger.debug(

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

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

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

658 _summarize_ast(ast),

659 )

660 return ast

661

662

663def parse_tree(query: str):

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

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

666 t0 = time.perf_counter()

667 tree = _get_lark_parser().parse(query)

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

669 return tree