PostgreSqlFormatter.hpp

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include "../SqlQueryFormatter.hpp"
#include "SQLiteFormatter.hpp"
 
#include <reflection-cpp/reflection.hpp>
 
#include <algorithm>
#include <format>
#include <string>
#include <string_view>
 
namespace Lightweight
{
 
class PostgreSqlFormatter final: public SQLiteQueryFormatter
{
  public:
    using SQLiteQueryFormatter::CreateTable;
 
    [[nodiscard]] bool RequiresTableRebuildForSchemaChange() const noexcept override
    {
        return false;
    }
 
    [[nodiscard]] bool IdentifiersAreCaseInsensitive() const noexcept override
    {
        return false;
    }
 
    [[nodiscard]] bool HasNativeGuidType() const noexcept override
    {
        return true;
    }
 
    [[nodiscard]] bool ComparesFixedCharWithPadSemantics() const noexcept override
    {
        return true;
    }
 
    [[nodiscard]] bool BoundsBulkLoadTransactionSize() const noexcept override
    {
        return false;
    }
 
    [[nodiscard]] std::string FoldUnquotedIdentifier(std::string_view identifier) const override
    {
        auto folded = std::string(identifier);
        std::ranges::transform(
            folded, folded.begin(), [](char c) { return (c >= 'A' && c <= 'Z') ? static_cast<char>(c - 'A' + 'a') : c; });
        return folded;
    }
 
    [[nodiscard]] bool SupportsSchemas() const noexcept override
    {
        return true;
    }
 
    [[nodiscard]] std::string CreateSchemaIfNotExists(std::string_view schema) const override
    {
        // CREATE SCHEMA IF NOT EXISTS checks the CREATE privilege on the database even when the schema exists.
        // The dollar-quote tag must not occur in the name, or the name would end the DO body.
        auto tag = std::string { "$lw$" };
        auto tagNumber = 0;
        while (schema.contains(tag))
            tag = std::format("$lw{}$", ++tagNumber);
        return std::format(R"(DO {0} BEGIN IF NOT EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = '{1}') THEN )"
                           R"(CREATE SCHEMA "{2}"; END IF; END {0})",
                           tag,
                           DoubleQuoteCharacter(schema, '\''),
                           DoubleQuoteCharacter(schema, '"'));
    }
 
    [[nodiscard]] StringList DropTable(std::string_view schemaName,
                                       std::string_view const& tableName,
                                       bool ifExists = false,
                                       bool cascade = false) const override
    {
        std::string sql = ifExists ? std::format("DROP TABLE IF EXISTS {}", FormatTableName(schemaName, tableName))
                                   : std::format("DROP TABLE {}", FormatTableName(schemaName, tableName));
        if (cascade)
            sql += " CASCADE";
        sql += ";";
        return { sql };
    }
 
    [[nodiscard]] std::string BinaryLiteral(std::span<uint8_t const> data) const override
    {
        std::string result;
        result.reserve((data.size() * 2) + 4);
        result += "'\\x";
        for (uint8_t byte: data)
            result += std::format("{:02X}", byte);
        result += "'";
        return result;
    }
 
    [[nodiscard]] std::string QueryLastInsertId(std::string_view /*tableName*/) const override
    {
        // NB: Find a better way to do this on the given table.
        // In our case it works, because we're expected to call this right after an insert.
        // But a race condition may still happen if another client inserts a row at the same time too.
        return std::format("SELECT lastval();");
    }
 
    [[nodiscard]] std::string_view DateFunction() const noexcept override
    {
        return "CURRENT_DATE";
    }
 
    /// PostgreSQL session default schema is selected via `search_path`. We
    /// pin the migration tool's chosen schema first and fall back to `public`
    /// so unqualified references to built-ins keep resolving. Schema names
    /// must be pre-validated by the caller (see
    /// `MigrationManager::SetDefaultSchema`).
    [[nodiscard]] std::string SetDefaultSchemaStatement(std::string_view schema) const override
    {
        if (schema.empty())
            return {};
        return std::format(R"(SET search_path TO "{}", public)", schema);
    }
 
    [[nodiscard]] std::string BuildColumnDefinition(SqlColumnDeclaration const& column) const override
    {
        std::stringstream sqlQueryString;
 
        sqlQueryString << '"' << column.name << "\" ";
 
        // Detect PostgreSQL auto-increment columns by checking for nextval() in default value.
        // This handles restore of backed-up tables where SERIAL columns have their default
        // value captured as nextval('"TableName_id_seq"'::regclass).
        bool const isAutoIncrementViaDefault = column.defaultValue.contains("nextval(");
        bool const isAutoIncrement = column.primaryKey == SqlPrimaryKeyType::AUTO_INCREMENT || isAutoIncrementViaDefault;
 
        if (isAutoIncrement)
            sqlQueryString << SerialColumnType(column.type);
        else
            sqlQueryString << ColumnType(column.type);
 
        if (column.required)
            sqlQueryString << " NOT NULL";
 
        // Only add inline PRIMARY KEY for explicitly marked AUTO_INCREMENT columns.
        // For columns detected via nextval() default, the table-level PRIMARY KEY constraint
        // will handle it to avoid "multiple primary keys" error.
        if (column.primaryKey == SqlPrimaryKeyType::AUTO_INCREMENT)
            sqlQueryString << " PRIMARY KEY";
        else if (column.primaryKey == SqlPrimaryKeyType::NONE && !column.index && column.unique)
            sqlQueryString << " UNIQUE";
 
        // Don't output default value for auto-increment columns as SERIAL handles it
        if (!column.defaultValue.empty() && !isAutoIncrement)
            sqlQueryString << " DEFAULT " << column.defaultValue;
 
        return sqlQueryString.str();
    }
 
    /// Maps a declared integer column type onto the matching PostgreSQL auto-increment pseudo-type.
    ///
    /// PostgreSQL picks the underlying integer width from the serial spelling, so the width must
    /// follow the declared type. Emitting a plain `SERIAL` for a `Bigint` key would silently create
    /// a 32-bit `integer` column that caps at 2^31-1 and no longer matches `BIGINT` foreign keys
    /// referencing it.
    ///
    /// @note `Smallint` and `Tinyint` map onto `SMALLSERIAL`, whose underlying `smallint` caps at
    ///       32767 — declaring a narrow auto-increment key buys a correspondingly narrow key space.
    /// @note A declared type with no integer equivalent at all (`Guid`, `Varchar`, `Decimal`, ...)
    ///       falls through to `SERIAL`, i.e. it is silently coerced to a 32-bit integer column.
    ///       PostgreSQL accepts that DDL, so the mismatch only surfaces later when reading the
    ///       column back into the declared C++ type. Auto-increment is only meaningful on an
    ///       integer key; prefer rejecting such records at the call site.
    ///
    /// @param type The declared column type.
    /// @return The serial pseudo-type to emit; `SERIAL` for any type without a narrower or wider
    ///         serial equivalent.
    [[nodiscard]] static std::string_view SerialColumnType(SqlColumnTypeDefinition const& type)
    {
        using namespace SqlColumnTypeDefinitions;
 
        return std::visit(detail::overloaded {
                              [](Bigint const&) -> std::string_view { return "BIGSERIAL"; },
                              [](Smallint const&) -> std::string_view { return "SMALLSERIAL"; },
                              // NB: PostgreSQL has no 1-byte integer type, so the narrowest
                              // available serial is used here, matching ColumnType(Tinyint).
                              [](Tinyint const&) -> std::string_view { return "SMALLSERIAL"; },
                              [](auto const&) -> std::string_view { return "SERIAL"; },
                          },
                          type);
    }
 
    /// The widest `VARCHAR(n)` PostgreSQL accepts. A wider column, such as a SQL Server `VARCHAR(MAX)` or `TEXT`
    /// column read from an archive, can only be `TEXT` here.
    static constexpr std::size_t MaxVarcharLength = 10'485'760;
 
    [[nodiscard]] std::string ColumnType(SqlColumnTypeDefinition const& type) const override
    {
        using namespace SqlColumnTypeDefinitions;
 
        // PostgreSQL stores all strings as UTF-8
        return std::visit(detail::overloaded {
                              [](Bigint const&) -> std::string { return "BIGINT"; },
                              [](Binary const& type) -> std::string { return std::format("BYTEA", type.size); },
                              [](Bool const&) -> std::string { return "BOOLEAN"; },
                              [](Char const& type) -> std::string { return std::format("CHAR({})", type.size); },
                              [](Date const&) -> std::string { return "DATE"; },
                              [](DateTime const&) -> std::string { return "TIMESTAMP"; },
                              [](Decimal const& type) -> std::string {
                                  return std::format("DECIMAL({}, {})", type.precision, type.scale);
                              },
                              [](Guid const&) -> std::string { return "UUID"; },
                              [](Integer const&) -> std::string { return "INTEGER"; },
                              [](NChar const& type) -> std::string { return std::format("CHAR({})", type.size); },
                              [](NVarchar const& type) -> std::string {
                                  if (type.size == 0 || type.size > MaxVarcharLength)
                                      return "TEXT";
                                  return std::format("VARCHAR({})", type.size);
                              },
                              [](Real const& type) -> std::string {
                                  // PostgreSQL REAL is float4; a Real with precision > 24 (e.g. an
                                  // introspected double-precision/float(53) column) must round-trip
                                  // as DOUBLE PRECISION or restore silently narrows it to float32.
                                  return type.precision > 24 ? "DOUBLE PRECISION" : "REAL";
                              },
                              [](Smallint const&) -> std::string { return "SMALLINT"; },
                              [](Text const&) -> std::string { return "TEXT"; },
                              [](Time const&) -> std::string { return "TIME"; },
                              [](Timestamp const&) -> std::string { return "TIMESTAMP"; },
                              // NB: PostgreSQL doesn't have a TINYINT type, but it does have a SMALLINT type.
                              [](Tinyint const&) -> std::string { return "SMALLINT"; },
                              [](VarBinary const& /*type*/) -> std::string { return std::format("BYTEA"); },
                              [](Varchar const& type) -> std::string {
                                  if (type.size == 0 || type.size > MaxVarcharLength)
                                      return "TEXT";
                                  return std::format("VARCHAR({})", type.size);
                              },
                          },
                          type);
    }
 
    // NOLINTNEXTLINE(readability-function-cognitive-complexity)
    [[nodiscard]] StringList AlterTable(std::string_view schemaName,
                                        std::string_view tableName,
                                        std::vector<SqlAlterTableCommand> const& commands) const override
    {
        std::stringstream sqlQueryString;
 
        int currentCommand = 0;
        for (SqlAlterTableCommand const& command: commands)
        {
            if (currentCommand > 0)
                sqlQueryString << '\n';
            ++currentCommand;
 
            using namespace SqlAlterTableCommands;
            sqlQueryString << std::visit(
                detail::overloaded {
                    [schemaName, tableName](RenameTable const& actualCommand) -> std::string {
                        return std::format(R"(ALTER TABLE {} RENAME TO "{}";)",
                                           FormatTableName(schemaName, tableName),
                                           actualCommand.newTableName);
                    },
                    [schemaName, tableName, this](AddColumn const& actualCommand) -> std::string {
                        return std::format(R"(ALTER TABLE {} ADD COLUMN "{}" {} {};)",
                                           FormatTableName(schemaName, tableName),
                                           actualCommand.columnName,
                                           ColumnType(actualCommand.columnType),
                                           actualCommand.nullable == SqlNullable::NotNull ? "NOT NULL" : "NULL");
                    },
                    [schemaName, tableName, this](AlterColumn const& actualCommand) -> std::string {
                        return std::format(
                            R"(ALTER TABLE {0} ALTER COLUMN "{1}" TYPE {2}, ALTER COLUMN "{1}" {3} NOT NULL;)",
                            FormatTableName(schemaName, tableName),
                            actualCommand.columnName,
                            ColumnType(actualCommand.columnType),
                            actualCommand.nullable == SqlNullable::NotNull ? "SET" : "DROP");
                    },
                    [schemaName, tableName](RenameColumn const& actualCommand) -> std::string {
                        return std::format(R"(ALTER TABLE {} RENAME COLUMN "{}" TO "{}";)",
                                           FormatTableName(schemaName, tableName),
                                           actualCommand.oldColumnName,
                                           actualCommand.newColumnName);
                    },
                    [schemaName, tableName](DropColumn const& actualCommand) -> std::string {
                        return std::format(R"(ALTER TABLE {} DROP COLUMN "{}";)",
                                           FormatTableName(schemaName, tableName),
                                           actualCommand.columnName);
                    },
                    [schemaName, tableName](AddIndex const& actualCommand) -> std::string {
                        using namespace std::string_view_literals;
                        auto const uniqueStr = actualCommand.unique ? "UNIQUE "sv : ""sv;
                        if (schemaName.empty())
                            return std::format(R"(CREATE {2}INDEX "{0}_{1}_index" ON "{0}" ("{1}");)",
                                               tableName,
                                               actualCommand.columnName,
                                               uniqueStr);
                        else
                            return std::format(R"(CREATE {3}INDEX "{0}_{1}_{2}_index" ON "{0}"."{1}" ("{2}");)",
                                               schemaName,
                                               tableName,
                                               actualCommand.columnName,
                                               uniqueStr);
                    },
                    [schemaName, tableName](DropIndex const& actualCommand) -> std::string {
                        if (schemaName.empty())
                            return std::format(R"(DROP INDEX "{0}_{1}_index";)", tableName, actualCommand.columnName);
                        else
                            return std::format(
                                R"(DROP INDEX "{0}_{1}_{2}_index";)", schemaName, tableName, actualCommand.columnName);
                    },
                    [schemaName, tableName](AddForeignKey const& actualCommand) -> std::string {
                        // Idempotent ADD CONSTRAINT — re-applying a migration must be a no-op.
                        // PostgreSQL has no native `IF NOT EXISTS` for `ADD CONSTRAINT`, so the
                        // guard is expressed via `DO $$ … EXCEPTION WHEN duplicate_object …`.
                        return std::format(
                            "DO $$ BEGIN ALTER TABLE {} ADD {}; EXCEPTION WHEN duplicate_object THEN NULL; END $$;",
                            FormatTableName(schemaName, tableName),
                            BuildForeignKeyConstraint(tableName, actualCommand.columnName, actualCommand.referencedColumn));
                    },
                    [schemaName, tableName](DropForeignKey const& actualCommand) -> std::string {
                        return std::format(R"(ALTER TABLE {} DROP CONSTRAINT "{}";)",
                                           FormatTableName(schemaName, tableName),
                                           BuildForeignKeyConstraintName(
                                               tableName, std::array { std::string_view { actualCommand.columnName } }));
                    },
                    [schemaName, tableName](AddCompositeForeignKey const& actualCommand) -> std::string {
                        std::stringstream ss;
                        ss << "ALTER TABLE " << FormatTableName(schemaName, tableName) << " ADD CONSTRAINT \""
                           << BuildForeignKeyConstraintName(tableName, actualCommand.columns) << "\" FOREIGN KEY (";
 
                        size_t i = 0;
                        for (auto const& col: actualCommand.columns)
                        {
                            if (i++ > 0)
                                ss << ", ";
                            ss << '"' << col << '"';
                        }
                        auto const referencedSchemaName = actualCommand.referencedSchemaName.empty()
                                                              ? schemaName
                                                              : std::string_view { actualCommand.referencedSchemaName };
                        ss << ") REFERENCES " << FormatTableName(referencedSchemaName, actualCommand.referencedTableName)
                           << " (";
 
                        i = 0;
                        for (auto const& col: actualCommand.referencedColumns)
                        {
                            if (i++ > 0)
                                ss << ", ";
                            ss << '"' << col << '"';
                        }
                        ss << ");";
                        return ss.str();
                    },
                    [schemaName, tableName, this](AddColumnIfNotExists const& actualCommand) -> std::string {
                        // PostgreSQL has native IF NOT EXISTS support for ADD COLUMN
                        return std::format(R"(ALTER TABLE {} ADD COLUMN IF NOT EXISTS "{}" {} {};)",
                                           FormatTableName(schemaName, tableName),
                                           actualCommand.columnName,
                                           ColumnType(actualCommand.columnType),
                                           actualCommand.nullable == SqlNullable::NotNull ? "NOT NULL" : "NULL");
                    },
                    [schemaName, tableName](DropColumnIfExists const& actualCommand) -> std::string {
                        // PostgreSQL has native IF EXISTS support for DROP COLUMN
                        return std::format(R"(ALTER TABLE {} DROP COLUMN IF EXISTS "{}";)",
                                           FormatTableName(schemaName, tableName),
                                           actualCommand.columnName);
                    },
                    [schemaName, tableName](DropIndexIfExists const& actualCommand) -> std::string {
                        // PostgreSQL has native IF EXISTS support for DROP INDEX
                        if (schemaName.empty())
                            return std::format(
                                R"(DROP INDEX IF EXISTS "{0}_{1}_index";)", tableName, actualCommand.columnName);
                        else
                            return std::format(R"(DROP INDEX IF EXISTS "{0}_{1}_{2}_index";)",
                                               schemaName,
                                               tableName,
                                               actualCommand.columnName);
                    },
                },
                command);
        }
 
        return { sqlQueryString.str() };
    }
 
    [[nodiscard]] std::string QueryServerVersion() const override
    {
        return "SELECT version()";
    }
 
    /// PostgreSQL has no SQLite settings: undoes the PRAGMAs inherited from the SQLite formatter.
    [[nodiscard]] StringList SqliteSettingsStatements(SqliteConnectionSettings const& /*settings*/) const override
    {
        return {};
    }
 
    /// PostgreSQL uses `pg_advisory_lock` / `pg_advisory_unlock`. Inline delegation
    /// keeps the vtable weak — see `SQLiteQueryFormatter::AdvisoryLockOps()` for
    /// the rationale.
    [[nodiscard]] SqlAdvisoryLockHandler const& AdvisoryLockOps() const override
    {
        return PostgreSqlAdvisoryLockOps();
    }
 
    /// PostgreSQL has a dedicated SQLSTATE for every recoverable condition.
    /// See `SQLiteQueryFormatter::AdvisoryLockOps()` for why this delegates.
    [[nodiscard]] SqlRetryClassifier const& RetryOps() const noexcept override
    {
        return PostgreSqlRetryOps();
    }
 
    /// The planner's tuple estimate of `pg_class`, which `ANALYZE` and autovacuum keep.
    [[nodiscard]] std::string_view TableRowCountEstimateQuery() const noexcept override
    {
        return "SELECT CAST(GREATEST(c.reltuples, 0) AS BIGINT) FROM pg_catalog.pg_class AS c "
               "JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace "
               "WHERE n.nspname = ? AND c.relname = ?";
    }
};
 
} // namespace Lightweight