progress

Author: psteinroe

crates/pgt_lexer/src/lexer.rs

@@ -111,6 +111,7 @@ impl<'a> Lexer<'a> {
                 pgt_tokenizer::TokenKind::Tilde => SyntaxKind::TILDE,
                 pgt_tokenizer::TokenKind::Question => SyntaxKind::QUESTION,
                 pgt_tokenizer::TokenKind::Colon => SyntaxKind::COLON,
+                pgt_tokenizer::TokenKind::DoubleColon => SyntaxKind::DOUBLE_COLON,
                 pgt_tokenizer::TokenKind::Eq => SyntaxKind::EQ,
                 pgt_tokenizer::TokenKind::Bang => SyntaxKind::BANG,
                 pgt_tokenizer::TokenKind::Lt => SyntaxKind::L_ANGLE,

crates/pgt_lexer_codegen/src/syntax_kind.rs

@@ -37,6 +37,7 @@ const PUNCT: &[(&str, &str)] = &[
     ("_", "UNDERSCORE"),
     (".", "DOT"),
     (":", "COLON"),
+    ("::", "DOUBLE_COLON"),
     ("=", "EQ"),
     ("!", "BANG"),
     ("-", "MINUS"),

crates/pgt_tokenizer/src/lib.rs

@@ -144,32 +144,37 @@ impl Cursor<'_> {
                 }
             }
             ':' => {
-                // Named parameters in psql with different substitution styles.
-                //
-                // https://www.postgresql.org/docs/current/app-psql.html#APP-PSQL-INTERPOLATION
-                match self.first() {
-                    '\'' => {
-                        // Named parameter with colon prefix and single quotes.
-                        self.bump();
-                        let terminated = self.single_quoted_string();
-                        let kind = NamedParamKind::ColonString { terminated };
-                        TokenKind::NamedParam { kind }
-                    }
-                    '"' => {
-                        // Named parameter with colon prefix and double quotes.
-                        self.bump();
-                        let terminated = self.double_quoted_string();
-                        let kind = NamedParamKind::ColonIdentifier { terminated };
-                        TokenKind::NamedParam { kind }
-                    }
-                    c if is_ident_start(c) => {
-                        // Named parameter with colon prefix.
-                        self.eat_while(is_ident_cont);
-                        TokenKind::NamedParam {
-                            kind: NamedParamKind::ColonRaw,
+                if self.first() == ':' {
+                    self.bump();
+                    TokenKind::DoubleColon
+                } else {
+                    // Named parameters in psql with different substitution styles.
+                    //
+                    // https://www.postgresql.org/docs/current/app-psql.html#APP-PSQL-INTERPOLATION
+                    match self.first() {
+                        '\'' => {
+                            // Named parameter with colon prefix and single quotes.
+                            self.bump();
+                            let terminated = self.single_quoted_string();
+                            let kind = NamedParamKind::ColonString { terminated };
+                            TokenKind::NamedParam { kind }
                         }
+                        '"' => {
+                            // Named parameter with colon prefix and double quotes.
+                            self.bump();
+                            let terminated = self.double_quoted_string();
+                            let kind = NamedParamKind::ColonIdentifier { terminated };
+                            TokenKind::NamedParam { kind }
+                        }
+                        c if is_ident_start(c) => {
+                            // Named parameter with colon prefix.
+                            self.eat_while(is_ident_cont);
+                            TokenKind::NamedParam {
+                                kind: NamedParamKind::ColonRaw,
+                            }
+                        }
+                        _ => TokenKind::Colon,
                     }
-                    _ => TokenKind::Colon,
                 }
             }
             // One-symbol tokens.
@@ -664,6 +669,23 @@ mod tests {
         assert_debug_snapshot!(result);
     }
 
+    #[test]
+    fn debug_simple_cast() {
+        let result = lex("::test");
+        assert_debug_snapshot!(result, @r###"
+        [
+            "::" @ DoubleColon,
+            "test" @ Ident,
+        ]
+        "###);
+    }
+
+    #[test]
+    fn named_param_colon_raw_vs_cast() {
+        let result = lex("select 1 from c where id::test = :id;");
+        assert_debug_snapshot!(result);
+    }
+
     #[test]
     fn named_param_colon_string() {
         let result = lex("select 1 from c where id = :'id';");

crates/pgt_tokenizer/src/snapshots/pgt_tokenizer__tests__named_param_colon_raw_vs_cast.snap

@@ -0,0 +1,25 @@
+---
+source: crates/pgt_tokenizer/src/lib.rs
+expression: result
+snapshot_kind: text
+---
+[
+    "select" @ Ident,
+    " " @ Space,
+    "1" @ Literal { kind: Int { base: Decimal, empty_int: false } },
+    " " @ Space,
+    "from" @ Ident,
+    " " @ Space,
+    "c" @ Ident,
+    " " @ Space,
+    "where" @ Ident,
+    " " @ Space,
+    "id" @ Ident,
+    "::" @ DoubleColon,
+    "test" @ Ident,
+    " " @ Space,
+    "=" @ Eq,
+    " " @ Space,
+    ":id" @ NamedParam { kind: ColonRaw },
+    ";" @ Semi,
+]

crates/pgt_tokenizer/src/token.rs

@@ -46,6 +46,8 @@ pub enum TokenKind {
     Minus,
     /// `:`
     Colon,
+    /// `::`
+    DoubleColon,
     /// `.`
     Dot,
     /// `=`

crates/pgt_workspace/src/workspace/server/pg_query.rs

@@ -142,6 +142,151 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_convert_to_positional_params_complex() {
+        let input = r#"CREATE OR REPLACE FUNCTION private.create_contacts_query(
+    v_and_filter jsonb DEFAULT NULL::jsonb,
+    v_include_filter jsonb DEFAULT NULL::jsonb,
+    v_exclude_filter jsonb DEFAULT NULL::jsonb,
+    v_require_marketing_opt_in boolean DEFAULT false,
+    v_require_transactional_opt_in boolean DEFAULT false,
+    v_include_contacts jsonb DEFAULT NULL::jsonb,
+    v_exclude_contacts jsonb DEFAULT NULL::jsonb,
+    v_channel_types public.channel_type[] DEFAULT NULL::public.channel_type[],
+    v_organisation_id uuid DEFAULT NULL::uuid,
+    v_include_blocked boolean DEFAULT false,
+    v_include_segment_ids uuid[] DEFAULT NULL::uuid[],
+    v_exclude_segment_ids uuid[] DEFAULT NULL::uuid[],
+    v_include_contact_list_ids uuid[] DEFAULT NULL::uuid[],
+    v_exclude_contact_list_ids uuid[] DEFAULT NULL::uuid[],
+    v_columns text[] DEFAULT NULL::text[],
+    -- expects an array of objects with `{"name": "column_name", "type": "column_type"}` format
+    -- used to include fields not present on the contact table such as `placeholder_values`
+    v_extra_fields jsonb DEFAULT NULL,
+    v_count_only boolean DEFAULT false,
+    -- below are fields that are only used in the UI
+    -- search is pushed down to the include subqueries
+    v_search text DEFAULT NULL,
+    -- order by is only allowed if v_limit is set
+    "v_order_by" "public"."column_sort"[] DEFAULT NULL::"public"."column_sort"[],
+    v_limit integer DEFAULT NULL,
+    v_offset integer DEFAULT NULL
+)
+ RETURNS text
+ LANGUAGE plpgsql
+ STABLE
+ SET search_path TO ''
+AS $function$
+declare
+    v_channel_types_specified public.channel_type[] := (
+        select array_agg(channel_type)
+        from unnest(coalesce(nullif(v_channel_types, '{}'), enum_range(null::public.channel_type))) channel_type
+    );
+    -- we need to include columns we filter on top-level
+    v_include_subqueries text[];
+    v_exclude_subqueries text[];
+    v_where_clauses text[];
+begin
+    if cardinality(v_columns) = 0 and v_count_only is not true then
+        raise exception using
+          message = 'No columns provided',
+          hint = 'Please pass v_columns',
+          errcode = 'INVIN';
+    end if;
+
+    if v_order_by is not null and v_limit is null then
+        raise exception using
+          message = 'v_order_by is only allowed if v_limit is set',
+          hint = 'Please pass v_limit',
+          errcode = 'INVIN';
+    end if;
+
+    if v_limit is not null and v_limit > 50 then
+        raise exception using
+          message = 'v_limit is too high',
+          hint = 'Please pass a v_limit of 50 or lower',
+          errcode = 'INVIN';
+    end if;
+
+    v_where_clauses := array_remove(
+        array[
+            -- is_blocked filter
+            (case when v_include_blocked is true then null else 'is_blocked is not true' end),
+            -- opt in filter
+            (
+                case
+                    when v_require_marketing_opt_in is true then
+                        '(' || (select string_agg(format('(%s is not null and %s_marketing_opt_in is true)', ct.type, ct.type), ' or ') from unnest(coalesce(v_channel_types, enum_range(null::public.channel_type))) ct(type)) || ')'
+                    when v_require_transactional_opt_in is true then
+                        '(' || (select string_agg(format('(%s is not null and %s_transactional_opt_in is not null)', channel_type, channel_type), ' or ') from unnest(coalesce(v_channel_types, enum_range(null::public.channel_type))) channel_type) || ')'
+                    when v_channel_types is not null and cardinality(v_channel_types) > 0 then
+                       '(' || (select string_agg(format('(%s is not null)', channel_type), ' or ') from unnest(v_channel_types) channel_type) || ')'
+                    else null
+                end
+            ),
+            -- organisation_id filter
+            format('"c"."organisation_id" = %L', v_organisation_id),
+            -- search filter
+            (case when nullif(v_search, '') is not null then
+                format('"c"."fts" @@ to_tsquery(''simple'', %L)', v_search)
+            end)
+
+        ],
+        null
+    );
+
+    -- select cols from public.contact
+    -- left join with contacts include / exclude contacts
+    --
+    -- where <where clauses>
+
+    return format(
+        $sql$
+            select %s
+            from public.contact
+            -- joins
+            -- where clause
+            %s
+            -- order by clause
+            %s
+            -- limit + offset
+            %s %s
+        $sql$,
+        (select string_agg(col, ', ') from unnest(v_columns) as col),
+        -- joins
+        -- where clause
+        (case when cardinality(v_where_clauses) > 0 then
+            'where ' || array_to_string(v_where_clauses, ' and ')
+        else
+            ''
+        end),
+        -- order by
+        case when v_order_by is not null and cardinality(v_order_by) > 0 then
+            format('order by %s', (
+                    select string_agg(format('%I.%I %s', 'contact', (x).column_name, case when (x).descending then 'desc' else 'asc' end), ', ')
+                    from unnest(v_order_by) as x
+                    where (x).column_name in (select column_name from information_schema.columns where table_schema = 'public' and table_name = 'contact')
+                )
+            )
+        when v_limit is not null then
+            'order by c.id' -- we need to order by something if we have a limit and offset, otherwise the results are not deterministic
+        else null
+        end,
+        -- limit
+        (case when v_limit is not null then format('limit %L', v_limit) end),
+        -- offset
+        (case when v_offset is not null then format('offset %L', v_offset) end)
+    );
+end
+$function$
+; "#;
+        let result = convert_to_positional_params(input);
+        assert_eq!(
+            result,
+            "select * from users where id = $1   and name = $2   and email = $3      ;"
+        );
+    }
+
     #[test]
     fn test_plpgsql_syntax_error() {
         let input = "