Skip to main content

squawk_ide/
semantic_tokens.rs

1use rowan::{NodeOrToken, TextRange};
2use salsa::Database as Db;
3use squawk_syntax::{
4    SyntaxElement, SyntaxKind,
5    ast::{self, AstNode},
6};
7
8use crate::db::{File, parse};
9use crate::file::InFile;
10use crate::goto_definition::goto_definition;
11use crate::location::LocationKind;
12
13fn highlight_param_mode(out: &mut SemanticTokenBuilder, mode: ast::ParamMode) {
14    match mode {
15        ast::ParamMode::ParamIn(param_in) => {
16            if let Some(token) = param_in.in_token() {
17                out.push_keyword(token.into());
18            }
19        }
20        ast::ParamMode::ParamInOut(param_in_out) => {
21            if let Some(token) = param_in_out.in_token() {
22                out.push_keyword(token.into());
23            }
24            if let Some(token) = param_in_out.inout_token() {
25                out.push_keyword(token.into());
26            }
27            if let Some(token) = param_in_out.out_token() {
28                out.push_keyword(token.into());
29            }
30        }
31        ast::ParamMode::ParamOut(param_out) => {
32            if let Some(token) = param_out.out_token() {
33                out.push_keyword(token.into());
34            }
35        }
36        ast::ParamMode::ParamVariadic(param_variadic) => {
37            if let Some(token) = param_variadic.variadic_token() {
38                out.push_keyword(token.into());
39            }
40        }
41    }
42}
43
44fn highlight_timezone(out: &mut SemanticTokenBuilder, timezone: ast::Timezone) {
45    match timezone {
46        ast::Timezone::WithTimezone(with_timezone) => {
47            if let Some(token) = with_timezone.with_token() {
48                out.push_type(token.into());
49            }
50            if let Some(token) = with_timezone.time_token() {
51                out.push_type(token.into());
52            }
53            if let Some(token) = with_timezone.zone_token() {
54                out.push_type(token.into());
55            }
56        }
57        ast::Timezone::WithoutTimezone(without_timezone) => {
58            if let Some(token) = without_timezone.without_token() {
59                out.push_type(token.into());
60            }
61            if let Some(token) = without_timezone.time_token() {
62                out.push_type(token.into());
63            }
64            if let Some(token) = without_timezone.zone_token() {
65                out.push_type(token.into());
66            }
67        }
68    }
69}
70
71fn highlight_type(out: &mut SemanticTokenBuilder, ty: ast::Type) {
72    match ty {
73        ast::Type::ArrayType(_) => (),
74        ast::Type::BitType(bit_type) => {
75            if let Some(token) = bit_type.setof_token() {
76                out.push_type(token.into());
77            }
78            if let Some(token) = bit_type.bit_token() {
79                out.push_type(token.into());
80            }
81        }
82        ast::Type::BitVaryingType(bit_varying_type) => {
83            if let Some(token) = bit_varying_type.setof_token() {
84                out.push_type(token.into());
85            }
86            if let Some(token) = bit_varying_type.bit_token() {
87                out.push_type(token.into());
88            }
89            if let Some(token) = bit_varying_type.varying_token() {
90                out.push_type(token.into());
91            }
92        }
93        ast::Type::VarcharType(varchar_type) => {
94            if let Some(token) = varchar_type.setof_token() {
95                out.push_type(token.into());
96            }
97            if let Some(token) = varchar_type.national_token() {
98                out.push_type(token.into());
99            }
100            if let Some(token) = varchar_type
101                .varchar_token()
102                .or_else(|| varchar_type.nchar_token())
103                .or_else(|| varchar_type.character_token())
104                .or_else(|| varchar_type.char_token())
105            {
106                out.push_type(token.into());
107            }
108            if let Some(token) = varchar_type.varying_token() {
109                out.push_type(token.into());
110            }
111        }
112        ast::Type::CharacterType(character_type) => {
113            if let Some(token) = character_type.setof_token() {
114                out.push_type(token.into());
115            }
116            if let Some(token) = character_type.national_token() {
117                out.push_type(token.into());
118            }
119            if let Some(token) = character_type
120                .nchar_token()
121                .or_else(|| character_type.character_token())
122                .or_else(|| character_type.char_token())
123            {
124                out.push_type(token.into());
125            }
126        }
127        ast::Type::DoubleType(double_type) => {
128            if let Some(token) = double_type.setof_token() {
129                out.push_type(token.into());
130            }
131            if let Some(token) = double_type.double_token() {
132                out.push_type(token.into());
133            }
134            if let Some(token) = double_type.precision_token() {
135                out.push_type(token.into());
136            }
137        }
138        ast::Type::IntervalType(interval_type) => {
139            if let Some(token) = interval_type.setof_token() {
140                out.push_type(token.into());
141            }
142            if let Some(token) = interval_type.interval_token() {
143                out.push_type(token.into());
144            }
145        }
146        ast::Type::PathType(path_type) => {
147            if let Some(token) = path_type.setof_token() {
148                out.push_type(token.into());
149            }
150        }
151        ast::Type::TimeType(time_type) => {
152            if let Some(token) = time_type.setof_token() {
153                out.push_type(token.into());
154            }
155            if let Some(token) = time_type.time_token() {
156                out.push_type(token.into());
157            }
158            if let Some(timezone) = time_type.timezone() {
159                highlight_timezone(out, timezone);
160            }
161        }
162        ast::Type::TimestampType(timestamp_type) => {
163            if let Some(token) = timestamp_type.setof_token() {
164                out.push_type(token.into());
165            }
166            if let Some(token) = timestamp_type.timestamp_token() {
167                out.push_type(token.into());
168            }
169            if let Some(timezone) = timestamp_type.timezone() {
170                highlight_timezone(out, timezone);
171            }
172        }
173    }
174}
175
176/// A semantic token with its position and classification.
177#[derive(Debug, Clone, PartialEq, Eq)]
178pub struct SemanticToken {
179    pub range: TextRange,
180    pub token_type: SemanticTokenType,
181    pub modifiers: Option<SemanticTokenModifier>,
182}
183
184#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
185#[repr(u8)]
186pub enum SemanticTokenModifier {
187    Definition = 0,
188    Readonly,
189    Documentation,
190}
191
192/// Semantic token types supported by the language server.
193#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
194pub enum SemanticTokenType {
195    Keyword,
196    String,
197    Bool,
198    Number,
199    Function,
200    Operator,
201    Punctuation,
202    Name,
203    NameRef,
204    Comment,
205    Column,
206    Type,
207    Parameter,
208    PositionalParam,
209    PropertyGraph,
210    Table,
211    Schema,
212}
213
214impl TryFrom<LocationKind> for SemanticTokenType {
215    type Error = LocationKind;
216
217    fn try_from(kind: LocationKind) -> Result<Self, Self::Error> {
218        match kind {
219            LocationKind::Aggregate | LocationKind::Function | LocationKind::Procedure => {
220                Ok(SemanticTokenType::Function)
221            }
222            LocationKind::Column => Ok(SemanticTokenType::Column),
223            LocationKind::JsonPath | LocationKind::Label | LocationKind::Property => {
224                Ok(SemanticTokenType::Name)
225            }
226            LocationKind::NamedArgParameter => Ok(SemanticTokenType::Parameter),
227            LocationKind::Schema => Ok(SemanticTokenType::Schema),
228            LocationKind::PropertyGraph => Ok(SemanticTokenType::PropertyGraph),
229            LocationKind::ElementTable
230            | LocationKind::Sequence
231            | LocationKind::Table
232            | LocationKind::View => Ok(SemanticTokenType::Table),
233            LocationKind::Type => Ok(SemanticTokenType::Type),
234            LocationKind::AccessMethod
235            | LocationKind::CaseExpr
236            | LocationKind::Channel
237            | LocationKind::Collation
238            | LocationKind::CommitBegin
239            | LocationKind::CommitEnd
240            | LocationKind::Constraint
241            | LocationKind::Conversion
242            | LocationKind::Cursor
243            | LocationKind::Database
244            | LocationKind::EventTrigger
245            | LocationKind::Extension
246            | LocationKind::ForeignDataWrapper
247            | LocationKind::Index
248            | LocationKind::Language
249            | LocationKind::Operator
250            | LocationKind::OperatorClass
251            | LocationKind::OperatorFamily
252            | LocationKind::Policy
253            | LocationKind::PreparedStatement
254            | LocationKind::PreparedTransaction
255            | LocationKind::Publication
256            | LocationKind::Role
257            | LocationKind::Rule
258            | LocationKind::Savepoint
259            | LocationKind::Server
260            | LocationKind::Statistics
261            | LocationKind::Subscription
262            | LocationKind::Tablespace
263            | LocationKind::TextSearchConfiguration
264            | LocationKind::TextSearchDictionary
265            | LocationKind::TextSearchParser
266            | LocationKind::TextSearchTemplate
267            | LocationKind::Trigger
268            | LocationKind::Window => Err(kind),
269        }
270    }
271}
272
273fn token_type_for_node<T: AstNode>(db: &dyn Db, node: InFile<&T>) -> Option<SemanticTokenType> {
274    let offset = node.value.syntax().text_range().start();
275    let location = goto_definition(db, InFile::new(node.file_id, offset))
276        .into_iter()
277        .next()?;
278
279    SemanticTokenType::try_from(location.kind).ok()
280}
281
282#[derive(Default)]
283struct SemanticTokenBuilder {
284    tokens: Vec<SemanticToken>,
285}
286
287impl SemanticTokenBuilder {
288    fn build(mut self) -> Vec<SemanticToken> {
289        self.tokens
290            .sort_by_key(|token| (token.range.start(), token.range.end()));
291        self.tokens
292    }
293
294    fn push_keyword(&mut self, syntax_element: SyntaxElement) {
295        self.push_token(syntax_element, SemanticTokenType::Keyword);
296    }
297
298    fn push_type(&mut self, syntax_element: SyntaxElement) {
299        self.push_token(syntax_element, SemanticTokenType::Type);
300    }
301
302    fn push_token(&mut self, syntax_element: SyntaxElement, token_type: SemanticTokenType) {
303        self.tokens.push(SemanticToken {
304            range: syntax_element.text_range(),
305            token_type,
306            modifiers: None,
307        });
308    }
309}
310
311#[salsa::tracked]
312pub fn semantic_tokens(
313    db: &dyn Db,
314    file: File,
315    range_to_highlight: Option<TextRange>,
316) -> Vec<SemanticToken> {
317    let parse = parse(db, file);
318    let tree = parse.tree();
319    let root = tree.syntax();
320
321    // Determine the root based on the given range.
322    let (root, range_to_highlight) = {
323        let source_file = root;
324        match range_to_highlight {
325            Some(range) => {
326                let node = match source_file.covering_element(range) {
327                    NodeOrToken::Node(it) => it,
328                    NodeOrToken::Token(it) => it.parent().unwrap_or_else(|| source_file.clone()),
329                };
330                (node, range)
331            }
332            None => (source_file.clone(), source_file.text_range()),
333        }
334    };
335
336    let mut out = SemanticTokenBuilder::default();
337
338    // Taken from: https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/ide/src/syntax_highlighting.rs#L267C21-L267C21
339    let preorder = root.preorder_with_tokens();
340    for event in preorder {
341        use rowan::WalkEvent::{Enter, Leave};
342
343        let range = match &event {
344            Enter(it) | Leave(it) => it.text_range(),
345        };
346
347        // Element outside of the viewport, no need to highlight
348        if range_to_highlight.intersect(range).is_none() {
349            continue;
350        }
351
352        match event {
353            Enter(NodeOrToken::Node(node)) => {
354                if let Some(name) = ast::AnyName::cast(node.clone())
355                    && let Some(token_type) = token_type_for_node(db, InFile::new(file, &name))
356                {
357                    out.push_token(name.syntax().clone().into(), token_type);
358                }
359
360                if let Some(ty) = ast::Type::cast(node.clone()) {
361                    highlight_type(&mut out, ty);
362                }
363
364                if let Some(mode) = ast::ParamMode::cast(node.clone()) {
365                    highlight_param_mode(&mut out, mode);
366                }
367
368                // Cleanup various operators that the textmate grammar
369                // highlights spuriously. These are for the select cases that
370                // aren't easily handled in the textmate grammar.
371                if let Some(like_clause) = ast::LikeClause::cast(node.clone())
372                    && let Some(token) = like_clause.like_token()
373                {
374                    out.push_keyword(token.into());
375                }
376                if let Some(not_null_constraint) = ast::NotNullConstraint::cast(node.clone())
377                    && let Some(token) = not_null_constraint.not_token()
378                {
379                    out.push_keyword(token.into());
380                }
381                if let Some(partition_for_values_in) = ast::PartitionForValuesIn::cast(node.clone())
382                    && let Some(token) = partition_for_values_in.in_token()
383                {
384                    out.push_keyword(token.into());
385                }
386            }
387            Enter(NodeOrToken::Token(token)) => {
388                if token.kind() == SyntaxKind::WHITESPACE {
389                    continue;
390                }
391                if token.kind() == SyntaxKind::POSITIONAL_PARAM {
392                    out.push_token(token.into(), SemanticTokenType::PositionalParam);
393                }
394            }
395            Leave(_) => {}
396        }
397    }
398
399    out.build()
400}
401
402#[cfg(test)]
403mod test {
404    use crate::db::{Database, File};
405    use insta::assert_snapshot;
406    use std::fmt::Write;
407
408    #[must_use]
409    fn semantic_tokens(sql: &str) -> String {
410        let db = Database::default();
411        let file = File::new(&db, sql.to_string().into());
412        let tokens = super::semantic_tokens(&db, file, None);
413
414        let mut result = String::new();
415        for token in tokens {
416            let start: usize = token.range.start().into();
417            let end: usize = token.range.end().into();
418            let token_text = &sql[start..end];
419            // TODO: once we get modfifiers, we'll need to update this
420            let modifiers_text = "";
421            writeln!(
422                result,
423                "{:?} @ {}..{}: {:?}{}",
424                token_text, start, end, token.token_type, modifiers_text
425            )
426            .unwrap();
427        }
428        result
429    }
430
431    #[test]
432    fn create_function_misc_params() {
433        assert_snapshot!(semantic_tokens(
434            "
435create function add(
436  in a int = 1,
437  inout b text default 'x',
438  in out c varchar(10)[],
439  variadic d int[]
440) returns int
441as 'select $1 + $2'
442language sql;
443",
444        ), @r#"
445        "add" @ 17..20: Function
446        "in" @ 24..26: Keyword
447        "a" @ 27..28: Parameter
448        "int" @ 29..32: Type
449        "inout" @ 40..45: Keyword
450        "b" @ 46..47: Parameter
451        "text" @ 48..52: Type
452        "in" @ 68..70: Keyword
453        "out" @ 71..74: Keyword
454        "c" @ 75..76: Parameter
455        "varchar" @ 77..84: Type
456        "variadic" @ 94..102: Keyword
457        "d" @ 103..104: Parameter
458        "int" @ 105..108: Type
459        "int" @ 121..124: Type
460        "#);
461    }
462
463    #[test]
464    fn create_function_param_mode_type() {
465        assert_snapshot!(semantic_tokens(
466            "
467create function f(int8 in int8)
468returns void
469as '' language sql;
470",
471        ), @r#"
472        "f" @ 17..18: Function
473        "int8" @ 19..23: Parameter
474        "in" @ 24..26: Keyword
475        "int8" @ 27..31: Type
476        "void" @ 41..45: Type
477        "#);
478    }
479
480    #[test]
481    fn create_function_percent_type() {
482        assert_snapshot!(semantic_tokens(
483            "
484create function f(a t.c%type) 
485returns t.b%type 
486as '' language plpgsql;
487",
488        ), @r#"
489        "f" @ 17..18: Function
490        "a" @ 19..20: Parameter
491        "#);
492    }
493
494    #[test]
495    fn select_keywords() {
496        assert_snapshot!(semantic_tokens("
497select 1 and, 2 select;
498"), @r#"
499        "and" @ 10..13: Column
500        "select" @ 17..23: Column
501        "#)
502    }
503
504    #[test]
505    fn positional_param() {
506        assert_snapshot!(semantic_tokens("
507select $1, $2;
508"), @r#"
509        "$1" @ 8..10: PositionalParam
510        "$2" @ 12..14: PositionalParam
511        "#)
512    }
513
514    #[test]
515    fn insert_column_list() {
516        assert_snapshot!(semantic_tokens(
517            "
518create table products (product_no bigint, name text, price text);
519insert into products (product_no, name, price) values
520    (1, 'Cheese', 9.99),
521    (2, 'Bread', 1.99),
522    (3, 'Milk', 2.99);
523",
524        ), @r#"
525        "products" @ 14..22: Table
526        "product_no" @ 24..34: Column
527        "bigint" @ 35..41: Type
528        "name" @ 43..47: Column
529        "text" @ 48..52: Type
530        "price" @ 54..59: Column
531        "text" @ 60..64: Type
532        "products" @ 79..87: Table
533        "product_no" @ 89..99: Column
534        "name" @ 101..105: Column
535        "price" @ 107..112: Column
536        "#)
537    }
538
539    #[test]
540    fn from_alias_column_types() {
541        assert_snapshot!(semantic_tokens(
542            "
543select *
544from f as t(a int, b jsonb, c text, x int, ca char(5)[], ia int[][], r text);
545",
546        ), @r#"
547        "t" @ 20..21: Table
548        "a" @ 22..23: Column
549        "int" @ 24..27: Type
550        "b" @ 29..30: Column
551        "jsonb" @ 31..36: Type
552        "c" @ 38..39: Column
553        "text" @ 40..44: Type
554        "x" @ 46..47: Column
555        "int" @ 48..51: Type
556        "ca" @ 53..55: Column
557        "char" @ 56..60: Type
558        "ia" @ 67..69: Column
559        "int" @ 70..73: Type
560        "r" @ 79..80: Column
561        "text" @ 81..85: Type
562        "#);
563    }
564
565    #[test]
566    fn json_table_columns() {
567        assert_snapshot!(semantic_tokens(
568            "
569select *
570from my_films,
571json_table(
572  js,
573  '$.favorites[*]' columns (
574    id for ordinality,
575    kind text path '$.kind'
576  )
577) as jt;
578",
579        ), @r#"
580        "id" @ 76..78: Column
581        "kind" @ 99..103: Column
582        "text" @ 104..108: Type
583        "jt" @ 132..134: Table
584        "#);
585    }
586
587    #[test]
588    fn xml_table_columns() {
589        assert_snapshot!(semantic_tokens(
590            "
591select *
592from xmltable(
593  '/root/item'
594  passing xmlparse(document '<root><item id=\"1\"/></root>')
595  columns
596    row_num for ordinality,
597    item_id integer path '@id'
598);
599",
600        ), @r#"
601        "row_num" @ 113..120: Column
602        "item_id" @ 141..148: Column
603        "integer" @ 149..156: Type
604        "#);
605    }
606
607    #[test]
608    fn cast_types() {
609        assert_snapshot!(semantic_tokens(
610            "
611select '1'::jsonb, '2'::json, cast(1 as integer), cast(1 as int4[][]), cast(1 as varchar(10));
612",
613        ), @r#"
614        "jsonb" @ 13..18: Type
615        "json" @ 25..29: Type
616        "integer" @ 41..48: Type
617        "int4" @ 61..65: Type
618        "varchar" @ 82..89: Type
619        "#);
620    }
621
622    #[test]
623    fn cast_double() {
624        assert_snapshot!(semantic_tokens(
625            "
626select '1'::double precision;
627",
628        ), @r#"
629        "double" @ 13..19: Type
630        "precision" @ 20..29: Type
631        "#);
632    }
633
634    #[test]
635    fn cast_time_and_timestamp_time_zone() {
636        assert_snapshot!(semantic_tokens(
637            "
638select cast(1 as timestamp with time zone), cast(1 as timestamp without time zone), cast(1 as time with time zone), cast(1 as time without time zone);
639",
640        ), @r#"
641        "timestamp" @ 18..27: Type
642        "with" @ 28..32: Type
643        "time" @ 33..37: Type
644        "zone" @ 38..42: Type
645        "timestamp" @ 55..64: Type
646        "without" @ 65..72: Type
647        "time" @ 73..77: Type
648        "zone" @ 78..82: Type
649        "time" @ 95..99: Type
650        "with" @ 100..104: Type
651        "time" @ 105..109: Type
652        "zone" @ 110..114: Type
653        "time" @ 127..131: Type
654        "without" @ 132..139: Type
655        "time" @ 140..144: Type
656        "zone" @ 145..149: Type
657        "#);
658    }
659
660    #[test]
661    fn cast_national_character_varying_type() {
662        assert_snapshot!(semantic_tokens(
663            "
664select 'foo'::national character varying;
665",
666        ), @r#"
667        "national" @ 15..23: Type
668        "character" @ 24..33: Type
669        "varying" @ 34..41: Type
670        "#);
671    }
672
673    #[test]
674    fn create_function_returns_setof_type() {
675        assert_snapshot!(semantic_tokens(
676            "
677create function f() returns setof int
678as 'select 1'
679language sql;
680",
681        ), @r#"
682        "f" @ 17..18: Function
683        "setof" @ 29..34: Type
684        "int" @ 35..38: Type
685        "#);
686    }
687
688    #[test]
689    fn create_table_temporal_primary_key_column_types() {
690        assert_snapshot!(semantic_tokens(
691            "
692-- temporal_primary_key
693CREATE TABLE addresses (
694    id int8 generated BY DEFAULT AS IDENTITY,
695    valid_range tstzrange NOT NULL DEFAULT tstzrange(now(), 'infinity', '[)'),
696    recipient text NOT NULL,
697    PRIMARY KEY (id, valid_range WITHOUT OVERLAPS)
698);
699",
700        ), @r#"
701        "addresses" @ 38..47: Table
702        "id" @ 54..56: Column
703        "int8" @ 57..61: Type
704        "valid_range" @ 100..111: Column
705        "tstzrange" @ 112..121: Type
706        "NOT" @ 122..125: Keyword
707        "tstzrange" @ 139..148: Function
708        "now" @ 149..152: Function
709        "recipient" @ 179..188: Column
710        "text" @ 189..193: Type
711        "NOT" @ 194..197: Keyword
712        "id" @ 221..223: Column
713        "valid_range" @ 225..236: Column
714        "#);
715    }
716
717    #[test]
718    fn like_clause_keyword() {
719        assert_snapshot!(semantic_tokens(
720            "
721create table products(a text);
722create table test (
723  like products
724);
725",
726        ), @r#"
727        "products" @ 14..22: Table
728        "a" @ 23..24: Column
729        "text" @ 25..29: Type
730        "test" @ 45..49: Table
731        "like" @ 54..58: Keyword
732        "products" @ 59..67: Table
733        "#)
734    }
735
736    #[test]
737    fn partition_for_values_in_keywords() {
738        assert_snapshot!(semantic_tokens(
739            "
740create table t(a int);
741create table t_1 partition of t for values in (1);
742",
743        ), @r#"
744        "t" @ 14..15: Table
745        "a" @ 16..17: Column
746        "int" @ 18..21: Type
747        "t_1" @ 37..40: Table
748        "t" @ 54..55: Table
749        "in" @ 67..69: Keyword
750        "#)
751    }
752
753    #[test]
754    fn positional_param_and_cast_type() {
755        assert_snapshot!(semantic_tokens(
756            "
757select $2::jsonb;
758",
759        ), @r#"
760        "$2" @ 8..10: PositionalParam
761        "jsonb" @ 12..17: Type
762        "#);
763    }
764
765    #[test]
766    fn select_target_column() {
767        assert_snapshot!(semantic_tokens(
768            "
769create table t(a int, b text);
770select a, b from t;
771",
772        ), @r#"
773        "t" @ 14..15: Table
774        "a" @ 16..17: Column
775        "int" @ 18..21: Type
776        "b" @ 23..24: Column
777        "text" @ 25..29: Type
778        "a" @ 39..40: Column
779        "b" @ 42..43: Column
780        "t" @ 49..50: Table
781        "#);
782    }
783
784    #[test]
785    fn select_target_qualified_column() {
786        assert_snapshot!(semantic_tokens(
787            "
788create table t(a int);
789select t.a from t;
790",
791        ), @r#"
792        "t" @ 14..15: Table
793        "a" @ 16..17: Column
794        "int" @ 18..21: Type
795        "t" @ 31..32: Table
796        "a" @ 33..34: Column
797        "t" @ 40..41: Table
798        "#);
799    }
800
801    #[test]
802    fn select_target_function_call() {
803        assert_snapshot!(semantic_tokens(
804            "
805create function f() returns int as 'select 1' language sql;
806select f();
807",
808        ), @r#"
809        "f" @ 17..18: Function
810        "int" @ 29..32: Type
811        "f" @ 68..69: Function
812        "#);
813    }
814
815    #[test]
816    fn select_function_arg_and_qualified_column() {
817        assert_snapshot!(semantic_tokens(
818            "
819create table t(a int);
820create function b(t) returns int as 'select 1' language sql;
821select b(t), t.b from t;
822",
823        ), @r#"
824        "t" @ 14..15: Table
825        "a" @ 16..17: Column
826        "int" @ 18..21: Type
827        "b" @ 40..41: Function
828        "t" @ 42..43: Type
829        "int" @ 53..56: Type
830        "b" @ 92..93: Function
831        "t" @ 94..95: Table
832        "t" @ 98..99: Table
833        "b" @ 100..101: Function
834        "t" @ 107..108: Table
835        "#);
836    }
837
838    #[test]
839    fn policy_field_style_function_call() {
840        assert_snapshot!(semantic_tokens(
841            "
842create table t(c int);
843create function x(t) returns int as 'select 1' language sql;
844create policy p on t
845  with check (t.x > 0 and t.c > 0);
846",
847        ), @r#"
848        "t" @ 14..15: Table
849        "c" @ 16..17: Column
850        "int" @ 18..21: Type
851        "x" @ 40..41: Function
852        "t" @ 42..43: Type
853        "int" @ 53..56: Type
854        "t" @ 104..105: Table
855        "t" @ 120..121: Table
856        "x" @ 122..123: Function
857        "t" @ 132..133: Table
858        "c" @ 134..135: Column
859        "#);
860    }
861
862    #[test]
863    fn with_cte_name() {
864        assert_snapshot!(semantic_tokens(
865            "
866with t as (
867  select 1
868)
869select * from t;
870",
871        ), @r#"
872        "t" @ 6..7: Table
873        "t" @ 40..41: Table
874        "#);
875    }
876
877    #[test]
878    fn create_property_graph() {
879        assert_snapshot!(semantic_tokens(
880            "
881create property graph foo
882  vertex tables (bar key (a) no properties);
883",
884        ), @r#"
885        "foo" @ 23..26: PropertyGraph
886        "#);
887    }
888
889    #[test]
890    fn select_target_schema_qualified() {
891        assert_snapshot!(semantic_tokens(
892            "
893create schema s;
894create table s.t(a int);
895select s.t.a from s.t;
896",
897        ), @r#"
898        "s" @ 15..16: Schema
899        "s" @ 31..32: Schema
900        "t" @ 33..34: Table
901        "a" @ 35..36: Column
902        "int" @ 37..40: Type
903        "s" @ 50..51: Schema
904        "t" @ 52..53: Table
905        "a" @ 54..55: Column
906        "s" @ 61..62: Schema
907        "t" @ 63..64: Table
908        "#);
909    }
910}