diff --git a/packages/plpgsql-deparser/__tests__/hydrate.test.ts b/packages/plpgsql-deparser/__tests__/hydrate.test.ts index 8b8cbea7..5696b2cf 100644 --- a/packages/plpgsql-deparser/__tests__/hydrate.test.ts +++ b/packages/plpgsql-deparser/__tests__/hydrate.test.ts @@ -24,7 +24,7 @@ $$`; expect(result.stats.parsedExpressions).toBeGreaterThan(0); }); - it('should hydrate assignment expressions', () => { + it('should hydrate assignment expressions with :=', () => { const sql = `CREATE FUNCTION test_func() RETURNS integer LANGUAGE plpgsql AS $$ @@ -50,6 +50,36 @@ $$`; } }); + it('should hydrate assignment expressions with = (not just :=)', () => { + // PL/pgSQL allows both := and = for assignments + // This test ensures = assignments are properly hydrated with valueExpr + const sql = `CREATE FUNCTION test_func() RETURNS void +LANGUAGE plpgsql +AS $$ +DECLARE + x text; +BEGIN + x = "constructive-simple-secrets".get('a', 'b'); +END; +$$`; + + const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult; + const result = hydratePlpgsqlAst(parsed); + + expect(result.errors).toHaveLength(0); + expect(result.stats.assignmentExpressions).toBeGreaterThan(0); + + const assignExpr = findExprByKind(result.ast, 'assign'); + expect(assignExpr).toBeDefined(); + if (assignExpr && assignExpr.kind === 'assign') { + expect(assignExpr.target).toBe('x'); + expect(assignExpr.value).toContain('constructive-simple-secrets'); + // Critical: valueExpr must be populated for AST-based transformations + expect(assignExpr.valueExpr).toBeDefined(); + expect(assignExpr.valueExpr).not.toBeNull(); + } + }); + it('should hydrate IF condition expressions', () => { const sql = `CREATE FUNCTION test_func(p_val integer) RETURNS text LANGUAGE plpgsql diff --git a/packages/plpgsql-deparser/src/hydrate.ts b/packages/plpgsql-deparser/src/hydrate.ts index e1f66e50..29ddd7df 100644 --- a/packages/plpgsql-deparser/src/hydrate.ts +++ b/packages/plpgsql-deparser/src/hydrate.ts @@ -316,6 +316,7 @@ function splitAssignment(query: string): { target: string; value: string } | nul try { const tokens = scanSync(query); let assignIndex = -1; + let assignTokenText = ''; let parenDepth = 0; let bracketDepth = 0; @@ -327,8 +328,29 @@ function splitAssignment(query: string): { target: string; value: string } | nul else if (token.text === '[') bracketDepth++; else if (token.text === ']') bracketDepth--; + // Check for := first (preferred PL/pgSQL assignment operator) if (token.text === ':=' && parenDepth === 0 && bracketDepth === 0) { assignIndex = i; + assignTokenText = ':='; + break; + } + + // Also check for = (valid PL/pgSQL assignment operator) + // But avoid => (named parameter syntax) by checking the previous token + if (token.text === '=' && parenDepth === 0 && bracketDepth === 0) { + // Make sure this isn't part of => (named parameter) + // or comparison operators like >=, <=, <>, != + const prevToken = i > 0 ? tokens.tokens[i - 1] : null; + const prevText = prevToken?.text || ''; + + // Skip if previous token suggests this is not an assignment + if (prevText === '>' || prevText === '<' || prevText === '!' || prevText === ':') { + continue; + } + + // This looks like an assignment with = + assignIndex = i; + assignTokenText = '='; break; } } @@ -343,13 +365,33 @@ function splitAssignment(query: string): { target: string; value: string } | nul return { target, value }; } catch (err) { + // Fallback: try to find := first, then = const colonIndex = query.indexOf(':='); - if (colonIndex === -1) { - return null; + if (colonIndex !== -1) { + const target = query.substring(0, colonIndex).trim(); + const value = query.substring(colonIndex + 2).trim(); + return { target, value }; } - const target = query.substring(0, colonIndex).trim(); - const value = query.substring(colonIndex + 2).trim(); - return { target, value }; + + // Try to find = (but be careful about >=, <=, <>, !=, =>) + // Find the first = that's not part of a comparison operator + for (let i = 0; i < query.length; i++) { + if (query[i] === '=') { + const prev = i > 0 ? query[i - 1] : ''; + const next = i < query.length - 1 ? query[i + 1] : ''; + + // Skip comparison operators + if (prev === '>' || prev === '<' || prev === '!' || prev === ':' || next === '>') { + continue; + } + + const target = query.substring(0, i).trim(); + const value = query.substring(i + 1).trim(); + return { target, value }; + } + } + + return null; } }