private async Task <MultiShardDataReader> ExecAsync(MultiShardConnection conn, MultiShardCommand cmd) { cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults; cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn; return(await cmd.ExecuteReaderAsync(CommandBehavior.Default, CancellationToken.None)); }
private static async Task <int> ExecuteReader(MultiShardCommand command) { var rows = 0; using (var reader = await command.ExecuteReaderAsync()) { while (await reader.ReadAsync()) { rows++; } } return(rows); }
public void TestQueryShardsCommandCancellationHandler() { List <ShardLocation> cancelledShards = new List <ShardLocation>(); CancellationTokenSource cts = new CancellationTokenSource(); using (MultiShardCommand cmd = _shardConnection.CreateCommand()) { Barrier barrier = new Barrier(cmd.Connection.Shards.Count() + 1); // If the threads don't meet the barrier by this time, then give up and fail the test TimeSpan barrierTimeout = TimeSpan.FromSeconds(10); cmd.CommandText = "WAITFOR DELAY '00:01:00'"; cmd.CommandTimeoutPerShard = 12; cmd.ShardExecutionCanceled += (obj, args) => { cancelledShards.Add(args.ShardLocation); }; cmd.ShardExecutionBegan += (obj, args) => { // If ShardExecutionBegan were only signaled by one thread, // then this would hang forever. barrier.SignalAndWait(barrierTimeout); }; Task cmdTask = cmd.ExecuteReaderAsync(cts.Token); bool syncronized = barrier.SignalAndWait(barrierTimeout); Assert.IsTrue(syncronized); // Cancel the command once execution begins // Sleeps are bad but this is just to really make sure // sqlclient has had a chance to begin command execution // Will not effect the test outcome Thread.Sleep(TimeSpan.FromSeconds(1)); cts.Cancel(); // Validate that the task was cancelled AssertExtensions.WaitAndAssertThrows <TaskCanceledException>(cmdTask); // Validate that the cancellation event was fired for all shards List <ShardLocation> allShards = _shardConnection.GetShardConnections().Select(l => l.Item1).ToList(); CollectionAssert.AreEquivalent(allShards, cancelledShards, "Expected command canceled event to be fired for all shards!"); } }
private async Task<MultiShardDataReader> ExecAsync(MultiShardConnection conn, MultiShardCommand cmd) { cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults; cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn; return await cmd.ExecuteReaderAsync(CommandBehavior.Default, CancellationToken.None); }