using System; using System.Collections.Generic; using System.Globalization; using System.Text; using System.Threading; using System.Threading.Tasks; using Npgsql; namespace Jellyfin.DbMigrator; /// /// Writes rows to a PostgreSQL database using batched INSERT statements. /// public static class PostgresBulkWriter { /// /// The maximum number of rows per INSERT batch. /// private const int BatchSize = 500; /// /// Inserts all rows into the specified PostgreSQL table using batched INSERT statements. /// When is , logs what would be inserted without writing. /// /// An open . /// The name of the target PostgreSQL table. /// The rows to insert, as dictionaries mapping column name to value. /// When , skips actual writes. /// A token to cancel the operation. /// The number of rows that were inserted (or would have been inserted in dry-run mode). public static async Task WriteTableAsync( NpgsqlConnection connection, string tableName, IReadOnlyList> rows, bool isDryRun, CancellationToken cancellationToken = default) { ArgumentNullException.ThrowIfNull(connection); ArgumentException.ThrowIfNullOrWhiteSpace(tableName); TableNameValidator.EnsureSafe(tableName); ArgumentNullException.ThrowIfNull(rows); if (rows.Count == 0) { return 0L; } // Collect column names from the first row. var columns = new List(rows[0].Keys); if (isDryRun) { Console.WriteLine( $" [dry-run] Would insert {rows.Count} rows into \"{tableName}\" " + $"({string.Join(", ", columns)})."); return rows.Count; } long inserted = 0L; for (int offset = 0; offset < rows.Count; offset += BatchSize) { int end = Math.Min(offset + BatchSize, rows.Count); int batchCount = end - offset; var sql = BuildInsertSql(tableName, columns, batchCount); var cmd = connection.CreateCommand(); await using (cmd.ConfigureAwait(false)) { cmd.CommandText = sql; int paramIndex = 0; for (int rowIdx = offset; rowIdx < end; rowIdx++) { var row = rows[rowIdx]; foreach (var col in columns) { string paramName = $"p{paramIndex.ToString(CultureInfo.InvariantCulture)}"; row.TryGetValue(col, out object? val); cmd.Parameters.AddWithValue(paramName, val ?? DBNull.Value); paramIndex++; } } await cmd.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false); inserted += batchCount; } } return inserted; } /// /// Advances the PostgreSQL integer sequence for each table that contains an Id column, /// so that future auto-generated primary keys do not conflict with migrated data. /// /// An open . /// The names of the tables whose sequences should be advanced. /// When , logs the SQL without executing it. /// A token to cancel the operation. public static async Task AdvanceSequencesAsync( NpgsqlConnection connection, IEnumerable tableNames, bool isDryRun, CancellationToken cancellationToken = default) { ArgumentNullException.ThrowIfNull(connection); ArgumentNullException.ThrowIfNull(tableNames); foreach (var tableName in tableNames) { // Check if the table has an "Id" column. bool hasIdColumn = await TableHasColumnAsync( connection, tableName, "Id", cancellationToken).ConfigureAwait(false); if (!hasIdColumn) { continue; } string sql = $"SELECT setval(pg_get_serial_sequence('{tableName}', 'Id'), " + $"COALESCE((SELECT MAX(\"Id\") FROM \"{tableName}\"), 1))"; if (isDryRun) { Console.WriteLine($" [dry-run] Would advance sequence: {sql}"); continue; } var cmd = connection.CreateCommand(); await using (cmd.ConfigureAwait(false)) { cmd.CommandText = sql; try { await cmd.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false); } catch (Exception ex) { // Sequence may not exist for tables without serial PK — log and continue. Console.WriteLine( $" Warning: Could not advance sequence for \"{tableName}\": {ex.Message}"); } } } } /// /// Returns the number of rows currently in the specified PostgreSQL table. /// /// An open . /// The name of the table to count. /// A token to cancel the operation. /// The row count, or -1 if the table does not exist. public static async Task CountRowsAsync( NpgsqlConnection connection, string tableName, CancellationToken cancellationToken = default) { ArgumentNullException.ThrowIfNull(connection); ArgumentException.ThrowIfNullOrWhiteSpace(tableName); TableNameValidator.EnsureSafe(tableName); var cmd = connection.CreateCommand(); await using (cmd.ConfigureAwait(false)) { cmd.CommandText = $"SELECT COUNT(*) FROM \"{tableName}\""; try { var result = await cmd.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false); return result is long count ? count : Convert.ToInt64(result, CultureInfo.InvariantCulture); } catch (NpgsqlException) { return -1L; } } } /// /// Builds a parameterised bulk INSERT SQL statement for the given table, columns, and row count. /// /// The target table name. /// The ordered list of column names. /// The number of value-rows to include. /// A parameterised INSERT statement. private static string BuildInsertSql(string tableName, IReadOnlyList columns, int rowCount) { var sb = new StringBuilder(); sb.Append(CultureInfo.InvariantCulture, $"INSERT INTO \"{tableName}\" ("); for (int i = 0; i < columns.Count; i++) { if (i > 0) { sb.Append(", "); } sb.Append(CultureInfo.InvariantCulture, $"\"{columns[i]}\""); } sb.Append(") VALUES "); int paramIndex = 0; for (int row = 0; row < rowCount; row++) { if (row > 0) { sb.Append(", "); } sb.Append('('); for (int col = 0; col < columns.Count; col++) { if (col > 0) { sb.Append(", "); } sb.Append(CultureInfo.InvariantCulture, $"@p{paramIndex.ToString(CultureInfo.InvariantCulture)}"); paramIndex++; } sb.Append(')'); } sb.Append(" ON CONFLICT DO NOTHING"); return sb.ToString(); } /// /// Checks whether a given column exists in a PostgreSQL table. /// /// An open . /// The table name to check. /// The column name to look for. /// A token to cancel the operation. /// if the column exists; otherwise, . private static async Task TableHasColumnAsync( NpgsqlConnection connection, string tableName, string columnName, CancellationToken cancellationToken = default) { var cmd = connection.CreateCommand(); await using (cmd.ConfigureAwait(false)) { cmd.CommandText = "SELECT COUNT(*) FROM information_schema.columns " + "WHERE table_name = @table AND column_name = @col"; cmd.Parameters.AddWithValue("table", tableName); cmd.Parameters.AddWithValue("col", columnName); var result = await cmd.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false); long count = result is long l ? l : Convert.ToInt64(result, CultureInfo.InvariantCulture); return count > 0; } } }