Skip to content

Commit 0eb785d

Browse files
bartovalbartoval
committed
feat: add Registered_domain ast support
1 parent 613f392 commit 0eb785d

7 files changed

Lines changed: 282 additions & 3 deletions

File tree

‎src/ast/visitor/__tests__/commands.test.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,29 @@ test('can visit CHANGE_POINT command', () => {
139139
expect(list).toEqual(['CHANGE_POINT']);
140140
});
141141

142+
test('can visit REGISTERED_DOMAIN command', () => {
143+
const { ast } = EsqlQuery.fromSrc(`
144+
FROM index
145+
| REGISTERED_DOMAIN parts = host
146+
`);
147+
const visitor = new Visitor()
148+
.on('visitExpression', () => {
149+
return null;
150+
})
151+
.on('visitRegisteredDomainCommand', () => {
152+
return 'REGISTERED_DOMAIN';
153+
})
154+
.on('visitCommand', () => {
155+
return null;
156+
})
157+
.on('visitQuery', (ctx) => {
158+
return [...ctx.visitCommands()].flat();
159+
});
160+
const list = visitor.visitQuery(ast).flat().filter(Boolean);
161+
162+
expect(list).toEqual(['REGISTERED_DOMAIN']);
163+
});
164+
142165
test('can visit RERANK command', () => {
143166
const { ast } = EsqlQuery.fromSrc(`
144167
FROM movies

‎src/ast/visitor/contexts.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import type {
1919
ESQLAstItem,
2020
ESQLAstJoinCommand,
2121
ESQLAstQueryExpression,
22+
ESQLAstRegisteredDomainCommand,
2223
ESQLAstRerankCommand,
2324
ESQLColumn,
2425
ESQLCommandOption,
@@ -561,6 +562,12 @@ export class CompletionCommandVisitorContext<
561562
Data extends SharedData = SharedData,
562563
> extends CommandVisitorContext<Methods, Data, ESQLAstCommand> {}
563564

565+
// REGISTERED_DOMAIN <qualifiedName> = <primaryExpression>
566+
export class RegisteredDomainCommandVisitorContext<
567+
Methods extends VisitorMethods = VisitorMethods,
568+
Data extends SharedData = SharedData,
569+
> extends CommandVisitorContext<Methods, Data, ESQLAstRegisteredDomainCommand> {}
570+
564571
// SAMPLE <probability> [SEED <seed>]
565572
export class SampleCommandVisitorContext<
566573
Methods extends VisitorMethods = VisitorMethods,

‎src/ast/visitor/global_visitor_context.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import type {
1313
ESQLAstHeaderCommand,
1414
ESQLAstJoinCommand,
1515
ESQLAstQueryExpression,
16+
ESQLAstRegisteredDomainCommand,
1617
ESQLAstRerankCommand,
1718
ESQLColumn,
1819
ESQLFunction,
@@ -209,6 +210,14 @@ export class GlobalVisitorContext<
209210
input as any
210211
);
211212
}
213+
case 'registered_domain': {
214+
if (!this.methods.visitRegisteredDomainCommand) break;
215+
return this.visitRegisteredDomainCommand(
216+
parent,
217+
commandNode as ESQLAstRegisteredDomainCommand,
218+
input as any
219+
);
220+
}
212221
case 'sample': {
213222
if (!this.methods.visitSampleCommand) break;
214223
return this.visitSampleCommand(parent, commandNode, input as any);
@@ -464,6 +473,15 @@ export class GlobalVisitorContext<
464473
return this.visitWithSpecificContext('visitCompletionCommand', context, input);
465474
}
466475

476+
public visitRegisteredDomainCommand(
477+
parent: contexts.VisitorContext | null,
478+
node: ESQLAstRegisteredDomainCommand,
479+
input: types.VisitorInput<Methods, 'visitRegisteredDomainCommand'>
480+
): types.VisitorOutput<Methods, 'visitRegisteredDomainCommand'> {
481+
const context = new contexts.RegisteredDomainCommandVisitorContext(this, node, parent);
482+
return this.visitWithSpecificContext('visitRegisteredDomainCommand', context, input);
483+
}
484+
467485
public visitSampleCommand(
468486
parent: contexts.VisitorContext | null,
469487
node: ESQLAstCommand,

‎src/ast/visitor/types.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,8 @@ export type CommandVisitorInput<Methods extends VisitorMethods> = AnyToVoid<
110110
VisitorInput<Methods, 'visitMvExpandCommand'> &
111111
VisitorInput<Methods, 'visitJoinCommand'> &
112112
VisitorInput<Methods, 'visitRerankCommand'> &
113-
VisitorInput<Methods, 'visitChangePointCommand'>
113+
VisitorInput<Methods, 'visitChangePointCommand'> &
114+
VisitorInput<Methods, 'visitRegisteredDomainCommand'>
114115
>;
115116

116117
/**
@@ -141,7 +142,8 @@ export type CommandVisitorOutput<Methods extends VisitorMethods> =
141142
| VisitorOutput<Methods, 'visitJoinCommand'>
142143
| VisitorOutput<Methods, 'visitRerankCommand'>
143144
| VisitorOutput<Methods, 'visitChangePointCommand'>
144-
| VisitorOutput<Methods, 'visitCompletionCommand'>;
145+
| VisitorOutput<Methods, 'visitCompletionCommand'>
146+
| VisitorOutput<Methods, 'visitRegisteredDomainCommand'>;
145147

146148
/* eslint-disable @typescript-eslint/no-explicit-any */
147149
export interface VisitorMethods<
@@ -192,6 +194,11 @@ export interface VisitorMethods<
192194
any,
193195
any
194196
>;
197+
visitRegisteredDomainCommand?: Visitor<
198+
contexts.RegisteredDomainCommandVisitorContext<Visitors, Data>,
199+
any,
200+
any
201+
>;
195202
visitSampleCommand?: Visitor<contexts.SampleCommandVisitorContext<Visitors, Data>, any, any>;
196203
visitCommandOption?: Visitor<contexts.CommandOptionVisitorContext<Visitors, Data>, any, any>;
197204
visitFuseCommand?: Visitor<contexts.FuseCommandVisitorContext<Visitors, Data>, any, any>;

‎src/parser/__tests__/registered_domain.test.ts‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import { EsqlQuery } from '../../composer/query';
9+
import { Walker } from '../../ast/walker';
10+
import type { ESQLAstRegisteredDomainCommand, ESQLFunction } from '../../types';
11+
12+
describe('REGISTERED_DOMAIN', () => {
13+
describe('correctly formatted', () => {
14+
it('parses basic example', () => {
15+
const src = `FROM index | REGISTERED_DOMAIN parts = host`;
16+
const { ast, errors } = EsqlQuery.fromSrc(src);
17+
const registeredDomain = Walker.match(ast, {
18+
type: 'command',
19+
name: 'registered_domain',
20+
}) as ESQLAstRegisteredDomainCommand;
21+
22+
expect(errors.length).toBe(0);
23+
expect(registeredDomain).toMatchObject({
24+
type: 'command',
25+
name: 'registered_domain',
26+
incomplete: false,
27+
});
28+
expect(registeredDomain.targetField).toMatchObject({
29+
type: 'column',
30+
name: 'parts',
31+
});
32+
expect(registeredDomain.expression).toMatchObject({
33+
type: 'column',
34+
name: 'host',
35+
});
36+
expect(registeredDomain.args).toHaveLength(1);
37+
expect(registeredDomain.args[0]).toMatchObject({
38+
type: 'function',
39+
subtype: 'binary-expression',
40+
name: '=',
41+
});
42+
});
43+
44+
it('parses the assignment structure correctly', () => {
45+
const src = `FROM index | REGISTERED_DOMAIN parts = host`;
46+
const { ast } = EsqlQuery.fromSrc(src);
47+
const registeredDomain = ast.commands[1] as ESQLAstRegisteredDomainCommand;
48+
49+
const assignment = registeredDomain.args[0] as ESQLFunction;
50+
expect(assignment.args[0]).toMatchObject({
51+
type: 'column',
52+
name: 'parts',
53+
});
54+
expect(assignment.args[1]).toMatchObject({
55+
type: 'column',
56+
name: 'host',
57+
});
58+
});
59+
60+
it('parses with dotted qualified name', () => {
61+
const src = `FROM index | REGISTERED_DOMAIN my.prefix = fqdn`;
62+
const { ast, errors } = EsqlQuery.fromSrc(src);
63+
const registeredDomain = Walker.match(ast, {
64+
type: 'command',
65+
name: 'registered_domain',
66+
}) as ESQLAstRegisteredDomainCommand;
67+
68+
expect(errors.length).toBe(0);
69+
expect(registeredDomain.targetField).toMatchObject({
70+
type: 'column',
71+
});
72+
});
73+
74+
it('parses with function expression as source', () => {
75+
const src = `FROM index | REGISTERED_DOMAIN parts = CONCAT(host, ".com")`;
76+
const { ast, errors } = EsqlQuery.fromSrc(src);
77+
const registeredDomain = Walker.match(ast, {
78+
type: 'command',
79+
name: 'registered_domain',
80+
}) as ESQLAstRegisteredDomainCommand;
81+
82+
expect(errors.length).toBe(0);
83+
expect(registeredDomain.expression).toMatchObject({
84+
type: 'function',
85+
name: 'concat',
86+
});
87+
});
88+
89+
it('parses with quoted target prefix', () => {
90+
const src = 'FROM index | REGISTERED_DOMAIN `my-prefix` = fqdn';
91+
const { ast, errors } = EsqlQuery.fromSrc(src);
92+
const registeredDomain = ast.commands[1] as ESQLAstRegisteredDomainCommand;
93+
94+
expect(errors.length).toBe(0);
95+
expect(registeredDomain.targetField).toMatchObject({
96+
type: 'column',
97+
name: 'my-prefix',
98+
quoted: true,
99+
});
100+
});
101+
102+
it('parses with parameter expression as source', () => {
103+
const src = 'FROM index | REGISTERED_DOMAIN parts = ?fqdn_param';
104+
const { ast, errors } = EsqlQuery.fromSrc(src);
105+
const registeredDomain = ast.commands[1] as ESQLAstRegisteredDomainCommand;
106+
107+
expect(errors.length).toBe(0);
108+
expect(registeredDomain.expression).toMatchObject({
109+
type: 'literal',
110+
literalType: 'param',
111+
value: 'fqdn_param',
112+
});
113+
});
114+
});
115+
116+
describe('incorrectly formatted', () => {
117+
it('errors on just the command keyword', () => {
118+
const { ast, errors } = EsqlQuery.fromSrc(`FROM index | REGISTERED_DOMAIN`);
119+
const registeredDomain = ast.commands[1] as ESQLAstRegisteredDomainCommand;
120+
121+
expect(errors.length).toBeGreaterThan(0);
122+
expect(registeredDomain).toMatchObject({
123+
name: 'registered_domain',
124+
incomplete: true,
125+
});
126+
expect(registeredDomain.targetField).toMatchObject({
127+
type: 'column',
128+
name: '',
129+
incomplete: true,
130+
});
131+
expect(registeredDomain.expression).toBeUndefined();
132+
});
133+
134+
it('errors on missing assignment', () => {
135+
const { ast, errors } = EsqlQuery.fromSrc(`FROM index | REGISTERED_DOMAIN parts`);
136+
const registeredDomain = ast.commands[1] as ESQLAstRegisteredDomainCommand;
137+
138+
expect(errors.length).toBeGreaterThan(0);
139+
expect(registeredDomain).toMatchObject({
140+
name: 'registered_domain',
141+
incomplete: true,
142+
});
143+
expect(registeredDomain.targetField).toMatchObject({
144+
type: 'column',
145+
name: 'parts',
146+
});
147+
expect(registeredDomain.expression).toBeUndefined();
148+
});
149+
150+
it('errors on missing expression after assignment', () => {
151+
const { ast, errors } = EsqlQuery.fromSrc(`FROM index | REGISTERED_DOMAIN parts =`);
152+
const registeredDomain = ast.commands[1] as ESQLAstRegisteredDomainCommand;
153+
154+
expect(errors.length).toBeGreaterThan(0);
155+
expect(registeredDomain).toMatchObject({
156+
name: 'registered_domain',
157+
incomplete: true,
158+
});
159+
expect(registeredDomain.targetField).toMatchObject({
160+
type: 'column',
161+
name: 'parts',
162+
});
163+
expect(registeredDomain.expression).toMatchObject({
164+
type: 'unknown',
165+
incomplete: true,
166+
});
167+
});
168+
});
169+
});

‎src/parser/core/cst_to_ast_converter.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -477,6 +477,12 @@ export class CstToAstConverter {
477477
return this.fromCompletionCommand(completionCommandCtx);
478478
}
479479

480+
const registeredDomainCommandCtx = ctx.registeredDomainCommand();
481+
482+
if (registeredDomainCommandCtx) {
483+
return this.fromRegisteredDomainCommand(registeredDomainCommandCtx);
484+
}
485+
480486
const sampleCommandCtx = ctx.sampleCommand();
481487

482488
if (sampleCommandCtx) {
@@ -1427,6 +1433,49 @@ export class CstToAstConverter {
14271433
return command;
14281434
}
14291435

1436+
// -------------------------------------------------------- REGISTERED_DOMAIN
1437+
1438+
private fromRegisteredDomainCommand(
1439+
ctx: cst.RegisteredDomainCommandContext
1440+
): ast.ESQLAstRegisteredDomainCommand {
1441+
const command = this.createCommand<'registered_domain', ast.ESQLAstRegisteredDomainCommand>(
1442+
'registered_domain',
1443+
ctx
1444+
);
1445+
1446+
const qualifiedNameCtx = ctx.qualifiedName();
1447+
1448+
if (qualifiedNameCtx && ctx.ASSIGN()) {
1449+
const targetField = this.toColumn(qualifiedNameCtx);
1450+
command.targetField = targetField;
1451+
1452+
const expression = this.fromPrimaryExpressionStrict(ctx.primaryExpression());
1453+
command.expression = expression;
1454+
1455+
if (expression.incomplete || expression.type === 'unknown') {
1456+
command.incomplete = true;
1457+
}
1458+
1459+
const assignment = this.toFunction(
1460+
ctx.ASSIGN().getText(),
1461+
ctx,
1462+
undefined,
1463+
'binary-expression'
1464+
);
1465+
assignment.args.push(targetField, expression);
1466+
assignment.location = this.extendLocationToArgs(assignment);
1467+
command.args.push(assignment);
1468+
} else if (qualifiedNameCtx) {
1469+
const targetField = this.toColumn(qualifiedNameCtx);
1470+
1471+
command.targetField = targetField;
1472+
command.incomplete = true;
1473+
command.args.push(targetField);
1474+
}
1475+
1476+
return command;
1477+
}
1478+
14301479
// ------------------------------------------------------------------- SAMPLE
14311480

14321481
private fromSampleCommand(ctx: cst.SampleCommandContext): ast.ESQLCommand<'sample'> {

‎src/types.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ export type ESQLAstCommand =
1818
| ESQLAstRerankCommand
1919
| ESQLAstCompletionCommand
2020
| ESQLAstFuseCommand
21-
| ESQLAstForkCommand;
21+
| ESQLAstForkCommand
22+
| ESQLAstRegisteredDomainCommand;
2223

2324
export type ESQLAstAllCommands = ESQLAstCommand | ESQLAstHeaderCommand;
2425

@@ -156,6 +157,11 @@ export interface ESQLAstMmrCommand extends ESQLCommand<'mmr'> {
156157
namedParameters?: ESQLSingleAstItem;
157158
}
158159

160+
export interface ESQLAstRegisteredDomainCommand extends ESQLCommand<'registered_domain'> {
161+
targetField: ESQLColumn;
162+
expression?: ESQLAstExpression;
163+
}
164+
159165
/**
160166
* Represents a PROMQL command.
161167
*

0 commit comments

Comments
 (0)