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;
}
}
}