sql-utilities/src/Strata.SqlTools.PostgreSql/Statements/StatementReader.cs

using System.Text;
using Strata.SqlTools.SqlBreakdown.Classes;
using Strata.SqlTools.SqlBreakdown.Enums.SQL;
using Strata.SqlTools.SqlBreakdown.Exceptions;
using SqlServerStatementReader = Strata.SqlTools.Statements.SqlServer.StatementReader;

namespace Strata.SqlTools.Statements.PostgreSql;

/// <summary>
/// PostgreSQL-specific tokenizer class that reads a string representation of a PostgreSQL SQL statement
/// and parses out each part as a token. Handles PostgreSQL's double-quoted identifiers, schema-qualified names,
/// single-quoted string literals, positional parameters, and PostgreSQL naming conventions.
/// </summary>
public class StatementReader : SqlServerStatementReader
{
    public StatementReader(string sqlStatement) : base(sqlStatement)
    {
    }

    /// <summary>
    /// Handles PostgreSQL-specific characters: double-quotes (") for delimited identifiers,
    /// single quotes (') for string literals, dollar sign ($) for positional parameters,
    /// colon (:) for named parameters, and at-sign (@) for named parameters.
    /// </summary>
    /// <returns>True if the character was handled; false otherwise.</returns>
    /// <summary>
    /// Attempts to handle additional PostgreSQL-specific characters that the base reader doesn't handle.
    /// </summary>
    /// <returns>True if the character was handled; false otherwise.</returns>
#pragma warning disable S3776 // Cognitive Complexity - Refactoring this would reduce clarity
    protected override bool TryHandleAdditionalCharacter()
    {
        if (CurrentCharacter == '"')
        {
            // PostgreSQL uses double quotes for delimited identifiers (case-sensitive)
            MovePosition();
            var quotedIdentifier = GrabStringValue();
            _currentToken = new Token(TokenType.ColumnIdentifier, quotedIdentifier);
            if (CurrentCharacter != '"')
            {
                throw new InvalidSyntaxException(
                    $"Invalid syntax at position {Position}. Expected closing double quote.");
            }
            MovePosition();
            return true;
        }

        if (CurrentCharacter == '\'')
        {
            // PostgreSQL uses single quotes for string literals
            MovePosition();
            var stringLiteral = GrabStringLiteral();
            _currentToken = new Token(TokenType.String, stringLiteral);
            if (CurrentCharacter != '\'')
            {
                throw new InvalidSyntaxException(
                    $"Invalid syntax at position {Position}. Expected closing single quote.");
            }
            MovePosition();
            return true;
        }

        if (CurrentCharacter == '$')
        {
            // PostgreSQL positional parameters: $1, $2, etc.
            MovePosition();
            if (char.IsDigit(CurrentCharacter))
            {
                var paramNumber = GrabNumberValue();
                _currentToken = new Token(TokenType.Parameter, $"${paramNumber}");
                return true;
            }
            throw new InvalidSyntaxException(
                $"Invalid syntax at position {Position}. Expected digit after $.");
        }

        if (CurrentCharacter == ':')
        {
            // PostgreSQL colon-prefixed named parameters: :userId
            MovePosition();
            if (char.IsLetter(CurrentCharacter) || CurrentCharacter == '_')
            {
                var paramName = GrabStringValue();
                _currentToken = new Token(TokenType.Parameter, $":{paramName}");
                return true;
            }
            throw new InvalidSyntaxException(
                $"Invalid syntax at position {Position}. Expected identifier after :.");
        }

        if (CurrentCharacter == '@')
        {
            // PostgreSQL at-sign named parameters: @userId (also SQL Server compatible)
            MovePosition();
            if (char.IsLetter(CurrentCharacter) || CurrentCharacter == '_')
            {
                var paramName = GrabStringValue();
                _currentToken = new Token(TokenType.Parameter, $"@{paramName}");
                return true;
            }
            throw new InvalidSyntaxException(
                $"Invalid syntax at position {Position}. Expected identifier after @.");
        }

        if (CurrentCharacter == '=')
        {
            // =, => (PostgreSQL hstore + other operations)
            MovePosition();
            if (TryMatchTwoCharOperator('>', "=>")) return true;
            _currentToken = new Token(TokenType.Operator, "=");
            return true;
        }

        if (CurrentCharacter == '|')
        {
            // |, ||
            MovePosition();
            if (TryMatchTwoCharOperator('|', "||")) return true;
            _currentToken = new Token(TokenType.Operator, "|");
            return true;
        }

        if (CurrentCharacter == '<')
        {
            // <, <=, <>, <<
            MovePosition();
            if (TryMatchTwoCharOperator('=', "<=")) return true;
            if (TryMatchTwoCharOperator('>', "<>")) return true;
            if (TryMatchTwoCharOperator('<', "<<")) return true;
            _currentToken = new Token(TokenType.Operator, "<");
            return true;
        }

        if (CurrentCharacter == '>')
        {
            // >, >=, >>
            MovePosition();
            if (TryMatchTwoCharOperator('=', ">=")) return true;
            if (TryMatchTwoCharOperator('>', ">>")) return true;
            _currentToken = new Token(TokenType.Operator, ">");
            return true;
        }

        if (CurrentCharacter == '.')
        {
            // Handle .. range operator (used in arrays and ranges)
            // and single . for column qualification (table.column)
            if (Position + 1 < Length && _sqlStatement[Position + 1] == '.')
            {
                MovePosition();
                MovePosition();
                _currentToken = new Token(TokenType.Operator, "..");
                return true;
            }
            // Single . is used for column qualification (table.column)
            // Return it as an Operator token
            MovePosition();
            _currentToken = new Token(TokenType.Operator, ".");
            return true;
        }

        return false;
    }
#pragma warning restore S3776

    /// <summary>
    /// Handles PostgreSQL-specific identifier prefixes: underscores (_) can start identifiers.
    /// </summary>
    /// <returns>True if the character was handled; false otherwise.</returns>
    protected override bool TryHandleIdentifierPrefix()
    {
        if (CurrentCharacter == '_')
        {
            var underscoreIdentifier = GrabStringValue();
            _currentToken = new Token(TokenType.ColumnIdentifier, underscoreIdentifier);
            return true;
        }

        return false;
    }

    /// <summary>
    /// Grabs a string literal value between single quotes, handling PostgreSQL's escaped quotes ('').
    /// </summary>
    /// <returns>The string literal value without the surrounding quotes.</returns>
    private string GrabStringLiteral()
    {
        var stringValue = new StringBuilder();
        while (CurrentCharacter != '\'' && CurrentCharacter != char.MinValue)
        {
            stringValue.Append(CurrentCharacter);
            MovePosition();

            // Handle escaped single quotes ('')
            if (CurrentCharacter == '\'')
            {
                var nextPos = Position + 1;
                if (nextPos < Length && _sqlStatement[nextPos] == '\'')
                {
                    // Double single-quote is an escape
                    stringValue.Append('\'');
                    MovePosition(); // Skip first quote
                    MovePosition(); // Skip second quote
                }
            }
        }

        return stringValue.ToString();
    }

    /// <summary>
    /// If the position is currently sitting on <paramref name="nextChar"/>, advances past it,
    /// emits <paramref name="twoCharOperator"/> as the current Operator token, and returns
    /// <c>true</c>. Otherwise leaves position untouched and returns <c>false</c>. Used by the
    /// multi-character operator dispatch (e.g. <c>&lt;</c>/<c>&lt;=</c>/<c>&lt;&gt;</c>/<c>&lt;&lt;</c>).
    /// </summary>
    private bool TryMatchTwoCharOperator(char nextChar, string twoCharOperator)
    {
        if (CurrentCharacter != nextChar)
        {
            return false;
        }
        MovePosition();
        _currentToken = new Token(TokenType.Operator, twoCharOperator);
        return true;
    }
}