Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,8 @@ internal BulkCopySimpleResultSet()
_results = new List<Result>();
}

internal int Count => _results.Count;

internal Result this[int idx] => _results[idx];

// Callback function for the tdsparser
Expand Down Expand Up @@ -151,11 +153,12 @@ public SourceColumnMetadata(ValueMethod method, bool isSqlType, bool isDataFeed)
public readonly bool IsDataFeed;
}

// The initial query will return three tables.
// The initial query will return four result sets, but the column aliases result set will
// be empty when not required.
// Transaction count has only one value in one column and one row
// MetaData has n columns but no rows
// Collation has 4 columns and n rows
// Column aliases has 3 columns and n rows
// Column aliases has 2 columns and n rows

private const int MetaDataResultId = 1;

Expand Down Expand Up @@ -191,6 +194,7 @@ public SourceColumnMetadata(ValueMethod method, bool isSqlType, bool isDataFeed)

private SqlBulkCopyColumnMappingCollection _columnMappings;
private SqlBulkCopyColumnMappingCollection _localColumnMappings;
private bool _localColumnMappingsResolveAliases;

private SqlConnection _connection;
private SqlTransaction _internalTransaction;
Expand Down Expand Up @@ -244,6 +248,7 @@ private int RowNumber

// Metadata caching fields for CacheMetadata option
internal BulkCopySimpleResultSet CachedMetadata { get; private set; }
private bool _cachedMetadataResolveAliases;
// Per-operation clone of the destination table metadata, used when CacheMetadata is
// enabled so that column-pruning in AnalyzeTargetAndCreateUpdateBulkCommand does not
// mutate the cached BulkCopySimpleResultSet.
Expand Down Expand Up @@ -373,6 +378,7 @@ public string DestinationTableName
}

CachedMetadata = null;
_cachedMetadataResolveAliases = false;
_destinationTableName = value;
}
}
Expand Down Expand Up @@ -484,6 +490,45 @@ private string CreateInitialQuery()
string objectName = ADP.BuildMultiPartName(parts);
string escapedObjectName = SqlServerEscapeHelper.EscapeStringAsLiteral(objectName);
string catalogNameStringLiteral = CatalogName is null ? null : SqlServerEscapeHelper.EscapeStringAsLiteral(CatalogName);
string createColumnAliasesTableQuery = """

CREATE TABLE #Column_Aliases
(
[Canonical_Column_Name] SYSNAME,
[Canonical_Column_Id] INT,
[Aliased_Column_Name] SYSNAME
)
""";
string populateColumnAliasesQuery = _localColumnMappingsResolveAliases
? $"""

EXEC sp_executesql N'
INSERT INTO #Column_Aliases ([Canonical_Column_Name], [Canonical_Column_Id], [Aliased_Column_Name])
SELECT [name], [column_id], ''$to_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 8
UNION ALL
SELECT [name], [column_id], ''$from_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 5
UNION ALL
SELECT [name], [column_id], ''$edge_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 2 AND [name] LIKE ''$edge[_]id[_]%''
UNION ALL
SELECT [name], [column_id], ''$node_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 2 AND [name] LIKE ''$node[_]id[_]%''',
N'@Object_ID INT', @Object_ID = @Object_ID
Comment thread
benrr101 marked this conversation as resolved.
"""
: string.Empty;
string removeShadowedColumnAliasesQuery = _localColumnMappingsResolveAliases
? $"""

DELETE FROM #Column_Aliases
WHERE [Aliased_Column_Name] IN (SELECT [name] FROM {CatalogName}.[sys].[all_columns] WHERE [object_id] = @Object_ID)
"""
: string.Empty;
string selectColumnAliasesQuery = """

SELECT [Canonical_Column_Name], [Aliased_Column_Name]
FROM #Column_Aliases
ORDER BY [Canonical_Column_Id] ASC

DROP TABLE #Column_Aliases
""";
// Specify the column names explicitly. This is to ensure that we can map to hidden
// columns (e.g. columns in temporal tables.) If the target table doesn't exist,
// OBJECT_ID will return NULL and @Column_Names will remain non-null. The subsequent
Expand Down Expand Up @@ -543,13 +588,7 @@ private string CreateInitialQuery()
DECLARE @Column_Name_Query NVARCHAR(MAX);
DECLARE @Column_Names NVARCHAR(MAX) = NULL;
DECLARE @Has_Sys_All_Columns_Permissions INT = HAS_PERMS_BY_NAME('{catalogNameStringLiteral}.[sys].[all_columns]', 'OBJECT', 'SELECT');

CREATE TABLE #Column_Aliases
(
[Canonical_Column_Name] SYSNAME,
[Canonical_Column_Id] INT,
[Aliased_Column_Name] SYSNAME
)
{createColumnAliasesTableQuery}

IF CAST(SERVERPROPERTY('EngineEdition') AS INT) = 6
BEGIN
Expand All @@ -567,17 +606,7 @@ IF CAST(SERVERPROPERTY('EngineEdition') AS INT) = 6
IF EXISTS (SELECT TOP 1 * FROM {CatalogName}.[sys].[all_columns] WHERE [object_id] = OBJECT_ID('{catalogNameStringLiteral}.[sys].[all_columns]') AND [name] = 'graph_type')
BEGIN
SET @Column_Name_Query_FILTER = N'WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) NOT IN (1, 3, 4, 6, 7)';

EXEC sp_executesql N'
INSERT INTO #Column_Aliases ([Canonical_Column_Name], [Canonical_Column_Id], [Aliased_Column_Name])
SELECT [name], [column_id], ''$to_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 8
UNION ALL
SELECT [name], [column_id], ''$from_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 5
UNION ALL
SELECT [name], [column_id], ''$edge_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 2 AND [name] LIKE ''$edge[_]id[_]%''
UNION ALL
SELECT [name], [column_id], ''$node_id'' FROM {catalogNameStringLiteral}.[sys].[all_columns] WHERE [object_id] = @Object_ID AND COALESCE([graph_type], 0) = 2 AND [name] LIKE ''$node[_]id[_]%''',
N'@Object_ID INT', @Object_ID = @Object_ID
{populateColumnAliasesQuery}
END
ELSE
BEGIN
Expand All @@ -586,9 +615,7 @@ UNION ALL
SET @Column_Name_Query = @Column_Name_Query_SELECT + ' FROM {catalogNameStringLiteral}.[sys].[all_columns] ' + @Column_Name_Query_FILTER + ' ' + @Column_Name_Query_SORT + ';'

EXEC sp_executesql @Column_Name_Query, N'@Object_ID INT, @Column_Names NVARCHAR(MAX) OUTPUT', @Object_ID = @Object_ID, @Column_Names = @Column_Names OUTPUT;

DELETE FROM #Column_Aliases
WHERE [Aliased_Column_Name] IN (SELECT [name] FROM {CatalogName}.[sys].[all_columns] WHERE [object_id] = @Object_ID)
{removeShadowedColumnAliasesQuery}
END

SELECT @Column_Names = COALESCE(@Column_Names, '*');
Expand All @@ -598,12 +625,7 @@ WHERE [Aliased_Column_Name] IN (SELECT [name] FROM {CatalogName}.[sys].[all_colu
SET FMTONLY OFF;

EXEC {CatalogName}..{TableCollationsStoredProc} N'{SchemaName}.{TableName}';

SELECT [Canonical_Column_Name], [Aliased_Column_Name]
FROM #Column_Aliases
ORDER BY [Canonical_Column_Id] ASC

DROP TABLE #Column_Aliases
{selectColumnAliasesQuery}
""";
}

Expand All @@ -614,7 +636,7 @@ DROP TABLE #Column_Aliases
private Task<BulkCopySimpleResultSet> CreateAndExecuteInitialQueryAsync(out BulkCopySimpleResultSet result)
{
// Check if we have valid cached metadata for the current destination table
if (CachedMetadata != null)
if (IsCachedMetadataValid())
{
SqlClientEventSource.Log.TryTraceEvent("SqlBulkCopy.CreateAndExecuteInitialQueryAsync | Info | Using cached metadata for table '{0}'", _destinationTableName);
result = CachedMetadata;
Expand Down Expand Up @@ -655,11 +677,18 @@ private Task<BulkCopySimpleResultSet> CreateAndExecuteInitialQueryAsync(out Bulk
}
}

internal bool IsCachedMetadataValid()
{
return CachedMetadata != null
&& (!_localColumnMappingsResolveAliases || _cachedMetadataResolveAliases);
}

private void CacheMetadataIfEnabled(BulkCopySimpleResultSet result)
{
if (IsCopyOption(SqlBulkCopyOptions.CacheMetadata))
{
CachedMetadata = result;
_cachedMetadataResolveAliases = _localColumnMappingsResolveAliases;
SqlClientEventSource.Log.TryTraceEvent("SqlBulkCopy.CacheMetadataIfEnabled | Info | Cached metadata for table '{0}'", _destinationTableName);
}
}
Expand Down Expand Up @@ -1068,6 +1097,7 @@ private void WriteMetaData(BulkCopySimpleResultSet internalResults)
public void ClearCachedMetadata()
{
CachedMetadata = null;
_cachedMetadataResolveAliases = false;
SqlClientEventSource.Log.TryTraceEvent("SqlBulkCopy.ClearCachedMetadata | Info | Metadata cache cleared");
}

Expand All @@ -1092,6 +1122,7 @@ private void Dispose(bool disposing)
_columnMappings = null;
_parser = null;
CachedMetadata = null;
_cachedMetadataResolveAliases = false;
_operationMetaData = null;
try
{
Expand Down Expand Up @@ -1632,6 +1663,39 @@ private void AppendColumnNameAndTypeName(StringBuilder query, string columnName,
query.Append(typeName);
}

private void ResetLocalColumnMappings()
{
_localColumnMappings = null;
_localColumnMappingsResolveAliases = false;
}

private bool ShouldResolveColumnAliases()
Comment thread
benrr101 marked this conversation as resolved.
{
if (_localColumnMappings is null)
{
return false;
}

for (int i = 0; i < _localColumnMappings.Count; i++)
{
if (IsGraphColumnAlias(_localColumnMappings[i].DestinationColumn))
{
return true;
}
}

return false;
}

private bool IsGraphColumnAlias(string name)
{
string unquotedName = UnquotedName(name);
return string.Equals(unquotedName, "$node_id", StringComparison.OrdinalIgnoreCase)
|| string.Equals(unquotedName, "$edge_id", StringComparison.OrdinalIgnoreCase)
|| string.Equals(unquotedName, "$from_id", StringComparison.OrdinalIgnoreCase)
|| string.Equals(unquotedName, "$to_id", StringComparison.OrdinalIgnoreCase);
}

private string UnquotedName(string name)
{
if (string.IsNullOrEmpty(name))
Expand Down Expand Up @@ -2328,6 +2392,11 @@ private void WriteRowSourceToServerCommon(int columnCount)
if (_localColumnMappings.Count > 0)
{
_localColumnMappings.ValidateCollection();
foreach (SqlBulkCopyColumnMapping bulkCopyColumn in _localColumnMappings)
{
bulkCopyColumn.MappedDestinationColumn = null;
Comment thread
benrr101 marked this conversation as resolved.
}

foreach (SqlBulkCopyColumnMapping bulkCopyColumn in _localColumnMappings)
{
if (bulkCopyColumn._internalSourceColumnOrdinal == -1)
Expand All @@ -2343,6 +2412,8 @@ private void WriteRowSourceToServerCommon(int columnCount)
_localColumnMappings.CreateDefaultMapping(columnCount);
}

_localColumnMappingsResolveAliases = ShouldResolveColumnAliases();

// perf: If the user specified all column ordinals we do not need to get a schematable
if (unspecifiedColumnOrdinals)
{
Expand Down Expand Up @@ -3057,7 +3128,7 @@ private void WriteToServerInternalRestContinuedAsync(BulkCopySimpleResultSet int
// Bulk copy task is completed at this moment.
if (task.IsCanceled)
{
sqlBulkCopy._localColumnMappings = null;
sqlBulkCopy.ResetLocalColumnMappings();
try
{
sqlBulkCopy.CleanUpStateObject();
Expand All @@ -3069,11 +3140,12 @@ private void WriteToServerInternalRestContinuedAsync(BulkCopySimpleResultSet int
}
else if (task.Exception != null)
{
sqlBulkCopy.ResetLocalColumnMappings();
source.SetException(task.Exception.InnerException);
}
else
{
sqlBulkCopy._localColumnMappings = null;
sqlBulkCopy.ResetLocalColumnMappings();
try
{
sqlBulkCopy.CleanUpStateObject(isCancelRequested: false);
Expand All @@ -3098,7 +3170,7 @@ private void WriteToServerInternalRestContinuedAsync(BulkCopySimpleResultSet int
}
else
{
_localColumnMappings = null;
ResetLocalColumnMappings();

try
{
Expand All @@ -3117,7 +3189,7 @@ private void WriteToServerInternalRestContinuedAsync(BulkCopySimpleResultSet int
}
catch (Exception ex) when (ADP.IsCatchableExceptionType(ex))
{
_localColumnMappings = null;
ResetLocalColumnMappings();

try
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,8 @@ public void Test()
using (DbDataReader reader = srcCmd.ExecuteReader())
{
IDictionary stats;
long expectedIduCount = DataTestUtility.IsAzureSynapse || DataTestUtility.IsAtLeastSQL2017() ? 2 : 1;
long expectedSelectCount = DataTestUtility.IsAzureSynapse ? 4 : 13;
long expectedSelectRows = DataTestUtility.IsAzureSynapse ? 4 : 15;
long expectedTransactions = DataTestUtility.IsAzureSynapse || DataTestUtility.IsAtLeastSQL2017() ? 2 : 1;
using (SqlBulkCopy bulkcopy = new SqlBulkCopy(dstConn))
{
bulkcopy.DestinationTableName = dstTable;
Expand All @@ -69,12 +67,12 @@ public void Test()

DataTestUtility.AssertEqualsWithDescription((long)3, stats["BuffersReceived"], "Unexpected BuffersReceived value.");
DataTestUtility.AssertEqualsWithDescription((long)3, stats["BuffersSent"], "Unexpected BuffersSent value.");
DataTestUtility.AssertEqualsWithDescription(expectedIduCount, stats["IduCount"], "Unexpected IduCount value.");
DataTestUtility.AssertEqualsWithDescription((long)0, stats["IduCount"], "Unexpected IduCount value.");
Comment thread
mdaigle marked this conversation as resolved.
DataTestUtility.AssertEqualsWithDescription(expectedSelectCount, stats["SelectCount"], "Unexpected SelectCount value.");
DataTestUtility.AssertEqualsWithDescription((long)3, stats["ServerRoundtrips"], "Unexpected ServerRoundtrips value.");
DataTestUtility.AssertEqualsWithDescription(expectedSelectRows, stats["SelectRows"], "Unexpected SelectRows value.");
DataTestUtility.AssertEqualsWithDescription((long)2, stats["SumResultSets"], "Unexpected SumResultSets value.");
DataTestUtility.AssertEqualsWithDescription(expectedTransactions, stats["Transactions"], "Unexpected Transactions value.");
DataTestUtility.AssertEqualsWithDescription((long)0, stats["Transactions"], "Unexpected Transactions value.");
}
}
}
Expand Down
Loading
Loading