using System;

using System.Collections.Generic;

using System.Collections.ObjectModel;

using System.Data;

using System.Data.Common;

using System.Threading;

using System.Threading.Tasks;

using System.Transactions;

using Npgsql.BackendMessages;

using Npgsql.Internal;

using Npgsql.Internal.Postgres;

using Npgsql.PostgresTypes;

using Npgsql.Util;

using NpgsqlTypes;

namespace Npgsql.Schema;

sealed class DbColumnSchemaGenerator

readonly RowDescriptionMessage _rowDescription;

readonly NpgsqlConnection _connection;

readonly bool _fetchAdditionalInfo;

internal DbColumnSchemaGenerator(NpgsqlConnection connection, RowDescriptionMessage rowDescription, bool fetchAdditionalInfo)

}

#region Columns queries

static string GenerateColumnsQuery(Version pgVersion, string columnFieldFilter) =>

/// <summary>

/// Stripped-down version of <see cref="GenerateColumnsQuery"/>, mainly to support Amazon Redshift.

/// </summary>

static string GenerateOldColumnsQuery(string columnFieldFilter) =>

#endregion Column queries

internal async Task<ReadOnlyCollection<T>> GetColumnSchema<T>(bool async, CancellationToken cancellationToken = default) where T : DbColumn

{

// This is mainly for Amazon Redshift

var oldQueryMode = _connection.PostgreSqlVersion < new Version(8, 2);

var numFields = _rowDescription.Count;

var result = new List<T?>(numFields);

for (var i = 0; i < numFields; i++)

result.Add(null);

var populatedColumns = 0;

if (_fetchAdditionalInfo)

{

// We have two types of fields - those which correspond to actual database columns

// and those that don't (e.g. SELECT 8). For the former we load lots of info from

// the backend (if fetchAdditionalInfo is true), for the latter we only have the RowDescription

var filters = new List<string>();

for (var index = 0; index < _rowDescription.Count; index++)

{

var f = _rowDescription[index];

// Only column fields

if (f.TableOID != 0)

filters.Add($"(attr.attrelid={f.TableOID} AND attr.attnum={f.ColumnAttributeNumber})");

}

var columnFieldFilter = string.Join(" OR ", filters);

if (columnFieldFilter != string.Empty)

{

var query = oldQueryMode

? GenerateOldColumnsQuery(columnFieldFilter)

: GenerateColumnsQuery(_connection.PostgreSqlVersion, columnFieldFilter);

using var scope = new TransactionScope(

TransactionScopeOption.Suppress,

async ? TransactionScopeAsyncFlowOption.Enabled : TransactionScopeAsyncFlowOption.Suppress);

using var connection = (NpgsqlConnection)((ICloneable)_connection).Clone();

await connection.Open(async, cancellationToken).ConfigureAwait(false);

using var cmd = new NpgsqlCommand(query, connection);

var reader = await cmd.ExecuteReader(async, CommandBehavior.Default, cancellationToken).ConfigureAwait(false);

try

{

while (async ? await reader.ReadAsync(cancellationToken).ConfigureAwait(false) : reader.Read())

{

var column = LoadColumnDefinition(reader, _connection.Connector!.DatabaseInfo, oldQueryMode);

for (var ordinal = 0; ordinal < numFields; ordinal++)

{

var field = _rowDescription[ordinal];

if (field.TableOID == column.TableOID &&

field.ColumnAttributeNumber == column.ColumnAttributeNumber)

{

populatedColumns++;

if (column.ColumnOrdinal.HasValue)

column = column.Clone();

// The column's ordinal is with respect to the resultset, not its table

column.ColumnOrdinal = ordinal;

result[ordinal] = (T?)(object)column;

}

}

}

}

finally

{

if (async)

await reader.DisposeAsync().ConfigureAwait(false);

else

reader.Dispose();

}

}

}

// We had some fields which don't correspond to regular table columns (or fetchAdditionalInfo is false).

// Fill in whatever info we have from the RowDescription itself

for (var i = 0; i < numFields; i++)

{

var column = (NpgsqlDbColumn?)(object?)result[i];

var field = _rowDescription[i];

if (column is null)

{

column = SetUpNonColumnField(field);

column.ColumnOrdinal = i;

result[i] = (T?)(object)column;

populatedColumns++;

}

column.ColumnName = field.Name;

column.IsAliased = column.BaseColumnName is null ? default(bool?) : (column.BaseColumnName != column.ColumnName);

}

if (populatedColumns != numFields)

throw new NpgsqlException("Could not load all columns for the resultset");

return result.AsReadOnly()!;

}

NpgsqlDbColumn LoadColumnDefinition(NpgsqlDataReader reader, NpgsqlDatabaseInfo databaseInfo, bool oldQueryMode)

{

// We don't set ColumnName here. It should always contain the column alias rather than

// the table column name (i.e. in case of "SELECT foo AS foo_alias"). It will be set later.

var column = new NpgsqlDbColumn

{

AllowDBNull = !reader.GetBoolean(reader.GetOrdinal("attnotnull")),

BaseCatalogName = _connection.Database!,

BaseSchemaName = reader.GetString(reader.GetOrdinal("nspname")),

BaseServerName = _connection.Host!,

BaseTableName = reader.GetString(reader.GetOrdinal("relname")),

BaseColumnName = reader.GetString(reader.GetOrdinal("attname")),

ColumnAttributeNumber = reader.GetInt16(reader.GetOrdinal("attnum")),

IsKey = reader.GetBoolean(reader.GetOrdinal("isprimarykey")),

IsReadOnly = !reader.GetBoolean(reader.GetOrdinal("is_updatable")),

IsUnique = reader.GetBoolean(reader.GetOrdinal("isunique")),

TableOID = reader.GetFieldValue<uint>(reader.GetOrdinal("attrelid")),

TypeOID = reader.GetFieldValue<uint>(reader.GetOrdinal("typoid"))

};

column.PostgresType = databaseInfo.ByOID[column.TypeOID];

column.DataTypeName = column.PostgresType.DisplayName; // Facets do not get included

var defaultValueOrdinal = reader.GetOrdinal("default");

column.DefaultValue = reader.IsDBNull(defaultValueOrdinal) ? null : reader.GetString(defaultValueOrdinal);

column.IsIdentity = !oldQueryMode && reader.GetBoolean(reader.GetOrdinal("isidentity"));

// Use a heuristic to discover old SERIAL columns

column.IsAutoIncrement =

column.IsIdentity == true ||

column.DefaultValue != null && column.DefaultValue.StartsWith("nextval(", StringComparison.Ordinal);

ColumnPostConfig(column, reader.GetInt32(reader.GetOrdinal("typmod")));

return column;

}

NpgsqlDbColumn SetUpNonColumnField(FieldDescription field)

{

// ColumnName and BaseColumnName will be set later

var column = new NpgsqlDbColumn

{

BaseCatalogName = _connection.Database!,

BaseServerName = _connection.Host!,

IsReadOnly = true,

DataTypeName = field.PostgresType.DisplayName,

TypeOID = field.TypeOID,

TableOID = field.TableOID,

ColumnAttributeNumber = field.ColumnAttributeNumber,

PostgresType = field.PostgresType

};

ColumnPostConfig(column, field.TypeModifier);

return column;

}

/// <summary>

/// Performs some post-setup configuration that's common to both table columns and non-columns.

/// </summary>

void ColumnPostConfig(NpgsqlDbColumn column, int typeModifier)

{

var serializerOptions = _connection.Connector!.SerializerOptions;

// Call GetRepresentationalType to also handle domain types

// Because NpgsqlCommandBuilder relies on NpgsqlDbType for correct type mapping

// And otherwise we'll get NpgsqlDbType.Unknown

column.NpgsqlDbType = column.PostgresType.GetRepresentationalType().DataTypeName.ToNpgsqlDbType();

if (serializerOptions.GetTypeInfo(typeof(object), serializerOptions.ToCanonicalTypeId(column.PostgresType)) is { } typeInfo)

{

column.DataType = typeInfo.Type;

column.IsLong = column.PostgresType.DataTypeName == DataTypeNames.Bytea;

if (column.PostgresType is PostgresCompositeType)

column.UdtAssemblyQualifiedName = typeInfo.Type.AssemblyQualifiedName;

}

var facets = column.PostgresType.GetFacets(typeModifier);

if (facets.Size != null)

column.ColumnSize = facets.Size;

if (facets.Precision != null)

column.NumericPrecision = facets.Precision;

if (facets.Scale != null)

column.NumericScale = facets.Scale;

}