public void TestShardCommandSucceedHandler()
{
var shardConnections = CreateConnections(10, () => { });
ConcurrentDictionary <ShardLocation, bool> passedLocations = new ConcurrentDictionary <ShardLocation, bool>();
using (var conn = new MultiShardConnection(shardConnections))
{
Func <CancellationToken, MockSqlCommand, DbDataReader> executeReaderFunc = (token, cmd) =>
{
Thread.Sleep(TimeSpan.FromSeconds(2));
return(new MockSqlDataReader());
};
MockSqlCommand mockCmd = new MockSqlCommand();
mockCmd.ExecuteReaderFunc = executeReaderFunc;
mockCmd.CommandText = "Select 1";
using (MultiShardCommand cmd = MultiShardCommand.Create(conn, mockCmd, 10))
{
cmd.ShardExecutionSucceeded += new EventHandler <ShardExecutionEventArgs>((obj, eventArgs) =>
{
Assert.IsTrue(shardConnections.Select(x => x.Item1).Contains(eventArgs.ShardLocation), "The ShardLocation passed to the event handler does not exist in the set of passed in ShardLocations");
passedLocations[eventArgs.ShardLocation] = true;
});
CommandBehavior behavior = CommandBehavior.Default;
cmd.ExecuteReader(behavior);
}
}
Assert.AreEqual(shardConnections.Count, passedLocations.Count, "Not every ShardLocation had its corresponding event handler invoked.");
}
public void TestQueryShardsInvalidShardStateSync()
{
// Get a shard and close it's connection
var shardSqlConnections = _shardConnection.GetShardConnections();
shardSqlConnections[1].Item2.Close();
try
{
// Execute
using (MultiShardCommand cmd = _shardConnection.CreateCommand())
{
cmd.CommandText = "SELECT dbNameField, Test_int_Field, Test_bigint_Field FROM ConsistentShardedTable";
cmd.CommandType = CommandType.Text;
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
}
}
}
catch (Exception ex)
{
if (ex is AggregateException)
{
var aex = (AggregateException)ex;
Logger.Log("Exception encountered: " + ex.Message);
throw aex.InnerExceptions.FirstOrDefault((e) => e is InvalidOperationException);
}
throw;
}
}
public void TestQueryShardsAsync()
{
// Create new sharded connection so we can test the OpenAsync call as well.
//
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = "SELECT dbNameField, Test_int_Field, Test_bigint_Field FROM ConsistentShardedTable";
cmd.CommandType = CommandType.Text;
using (MultiShardDataReader sdr = ExecAsync(conn, cmd).Result)
{
int recordsRetrieved = 0;
while (sdr.Read())
{
recordsRetrieved++;
var dbNameField = sdr.GetString(0);
var testIntField = sdr.GetFieldValue <int>(1);
var testBigIntField = sdr.GetFieldValue <Int64>(2);
Logger.Log("RecordRetrieved: dbNameField: {0}, TestIntField: {1}, TestBigIntField: {2}, RecordCount: {3}",
dbNameField, testIntField, testBigIntField, recordsRetrieved);
}
Assert.AreEqual(recordsRetrieved, 9);
}
}
}
}
public void TestShardConnectionOpenException()
{
try
{
Action executeOnOpen = () =>
{
throw new InsufficientMemoryException();
};
var shardConnections = CreateConnections(10, executeOnOpen);
var mockCmd = new MockSqlCommand();
mockCmd.CommandText = "Select 1";
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 100))
{
cmd.CommandText = "select 1";
cmd.ExecuteReader();
}
}
}
catch (Exception ex)
{
if (ex is MultiShardAggregateException)
{
var maex = (MultiShardAggregateException)ex;
Logger.Log("Exception message: {0}.\n Exception tostring: {1}",
ex.Message, ex.ToString());
throw (maex.InnerException).InnerException;
}
throw;
}
}
public MultiShardAggregateException TestSelectFailure(string commandText, MultiShardExecutionPolicy policy)
{
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = commandText;
cmd.ExecutionPolicy = policy;
// ExecuteReader should fail
MultiShardAggregateException aggregateException =
AssertExtensions.AssertThrows <MultiShardAggregateException>(() => cmd.ExecuteReader());
// Sanity check the exceptions are the correct type
foreach (Exception e in aggregateException.InnerExceptions)
{
Assert.IsInstanceOfType(e, typeof(MultiShardException));
Assert.IsInstanceOfType(e.InnerException, typeof(SqlException));
}
// Return the exception so that the caller can do additional validation
return(aggregateException);
}
}
}
public void TestSelectNoRows(string commandText, MultiShardExecutionPolicy policy)
{
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = commandText;
cmd.ExecutionPolicy = policy;
// Read first
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
Assert.AreEqual(0, sdr.MultiShardExceptions.Count);
while (sdr.Read())
{
Assert.Fail("Should not have gotten any records.");
}
Assert.IsFalse(sdr.HasRows);
}
// HasRows first
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
Assert.AreEqual(0, sdr.MultiShardExceptions.Count);
Assert.IsFalse(sdr.HasRows);
while (sdr.Read())
{
Assert.Fail("Should not have gotten any records.");
}
}
}
}
}
public void TestFailedCommandWithConnectionCloseCmdBehavior()
{
Parallel.For(0, 100, i =>
{
try
{
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = "select * from table_does_not_exist";
cmd.CommandType = CommandType.Text;
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
while (sdr.Read())
{
}
}
}
}
}
catch (Exception ex)
{
Console.WriteLine("Encountered exception: {0} in iteration: {1}",
ex.ToString(), i);
}
finally
{
Console.WriteLine("Completed execution of iteration: {0}", i);
}
});
}
public void TestShardCommandCancellationHandler()
{
// Create connections to a few shards
var shardConnections = CreateConnections(10, () => { });
var mockCmd = new MockSqlCommand();
var cmdStartEvent = new ManualResetEvent(false);
mockCmd.ExecuteReaderFunc = (token, cmd) =>
{
while (true)
{
if (token == null)
{
break;
}
token.ThrowIfCancellationRequested();
Thread.Sleep(500);
cmdStartEvent.Set();
}
return(new MockSqlDataReader());
};
mockCmd.CommandText = "select 1";
ConcurrentDictionary <ShardLocation, bool> passedLocations = new ConcurrentDictionary <ShardLocation, bool>();
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 300))
{
cmd.ShardExecutionCanceled += new EventHandler <ShardExecutionEventArgs>((obj, eventArgs) =>
{
Assert.IsTrue(shardConnections.Select(x => x.Item1).Contains(eventArgs.ShardLocation),
"The ShardLocation passed to the event handler does not exist in the set of passed in ShardLocations");
passedLocations[eventArgs.ShardLocation] = true;
});
try
{
// start the Cancel on a separate thread
Task executeTask = Task.Run(() =>
{
cmdStartEvent.WaitOne();
cmd.Cancel();
});
cmd.ExecuteReader();
executeTask.Wait();
Assert.Fail("We should always be throwing an exception.");
}
catch (Exception ex)
{
Assert.IsTrue(ex is OperationCanceledException, "OperationCanceledException expected. Found {0}!", ex.ToString());
}
}
}
Assert.AreEqual(shardConnections.Count, passedLocations.Count, "Not every ShardLocation had its corresponding event handler invoked.");
}
/// <summary>
/// Executes the SQL command and returns the output in text format.
/// </summary>
private static string ExecuteCommand(MultiShardConnection conn, string commandText)
{
try
{
StringBuilder output = new StringBuilder();
output.AppendLine();
int rowsAffected = 0;
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = commandText;
cmd.CommandTimeout = s_commandLine.QueryTimeout;
cmd.CommandTimeoutPerShard = s_commandLine.QueryTimeout;
// Execute command and time with a stopwatch
Stopwatch stopwatch = Stopwatch.StartNew();
cmd.ExecutionPolicy = s_commandLine.ExecutionPolicy;
cmd.ExecutionOptions = s_commandLine.ExecutionOptions;
using (MultiShardDataReader reader = cmd.ExecuteReader(CommandBehavior.Default))
{
stopwatch.Stop();
// Get column names
IEnumerable <string> columnNames = GetColumnNames(reader).ToArray();
// Create table formatter
TableFormatter tableFormatter = new TableFormatter(columnNames.ToArray());
// Read results from db
while (reader.Read())
{
rowsAffected++;
// Add the row to the table formatter
object[] values = new object[reader.FieldCount];
reader.GetValues(values);
tableFormatter.AddRow(values);
}
// Write formatter output
output.AppendLine(tableFormatter.ToString());
}
output.AppendLine();
output.AppendFormat("({0} rows affected - {1:hh}:{1:mm}:{1:ss} elapsed)", rowsAffected, stopwatch.Elapsed);
output.AppendLine();
}
return(output.ToString());
}
catch (MultiShardAggregateException e)
{
return(e.ToString());
}
}
public void TestSchemaMismatchErrorPropagation()
{
// First we need to alter the schema on one of the shards - we'll choose the last one.
//
string origColName = "Test_bigint_Field";
string newColName = "ModifiedName";
MultiShardTestUtils.ChangeColumnNameOnShardedTable(2, origColName, newColName);
// Then create new sharded connection so we can test the error handling logic.
// We'll wrap this all in a try-catch-finally block so that we can change the schema back
// to what the other tests will expect it to be in the finally.
//
try
{
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
// Need to do a SELECT * in order to get the column name error as a schema mismatcherror. If we name it explicitly
// we will get a command execution error instead.
//
cmd.CommandText = "SELECT * FROM ConsistentShardedTable";
cmd.CommandType = CommandType.Text;
using (MultiShardDataReader sdr = ExecAsync(conn, cmd).Result)
{
// The number of errors we have depends on which shard executed first.
// So, we know it should be 1 OR 2.
//
Assert.IsTrue(
((sdr.MultiShardExceptions.Count == 1) || (sdr.MultiShardExceptions.Count == 2)),
string.Format("Expected 1 or 2 execution erros, but saw {0}", sdr.MultiShardExceptions.Count));
int recordsRetrieved = 0;
while (sdr.Read())
{
recordsRetrieved++;
var dbNameField = sdr.GetString(0);
}
// We should see 9 records less 3 for each one that had a schema error.
int expectedRecords = ((9 - (3 * sdr.MultiShardExceptions.Count)));
Assert.AreEqual(recordsRetrieved, expectedRecords);
}
}
}
}
finally
{
MultiShardTestUtils.ChangeColumnNameOnShardedTable(2, newColName, origColName);
}
}
// GET: Captain
public ActionResult Index()
{
Dictionary <int, int> votes = new Dictionary <int, int>();
//Perform a multi shard query to tally all the votes
// Get the shards to connect to
var shards = _shardMap.GetShards();
// Create the multi-shard connection
using (var conn = new MultiShardConnection(shards, Util.SqlHelper.GetCredentialsConnectionString()))
{
// Create a simple command
using (MultiShardCommand cmd = conn.CreateCommand())
{
// Because this query is grouped by CustomerID, which is sharded,
// we will not get duplicate rows.
cmd.CommandText = @"SELECT Captain, COUNT(*) AS Votes FROM [dbo].[Votes] GROUP BY [Captain]";
// Allow for partial results in case some shards do not respond in time
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
// Allow the entire command to take up to 30 seconds
cmd.CommandTimeout = 30;
// Execute the command.
// We do not need to specify retry logic because MultiShardDataReader will internally retry until the CommandTimeout expires.
using (MultiShardDataReader reader = cmd.ExecuteReader())
{
while (reader.Read())
{
votes.Add(reader.GetInt32(0), reader.GetInt32(1));
}
}
}
}
var vm = new ViewModels.Captain.IndexViewModel();
foreach (var captain in _captainContext.Captains)
{
int voteCount = 0;
votes.TryGetValue(captain.ID, out voteCount);
vm.Captains.Add(new ViewModels.Captain.IndexViewModelCaptain()
{
ID = captain.ID,
Name = captain.Name,
Votes = voteCount
});
}
return(View(vm));
}
// POST api/CustomLogin
public HttpResponseMessage Post(LoginRequest loginRequest)
{
Guid shardKey;
// SEND A QUERY TO ALL SHARD TO DETECT OUR SHARD!!!!
// SAVE companiesId to shardKey!
using (MultiShardConnection conn = new MultiShardConnection(WebApiConfig.ShardingObj.ShardMap.GetShards(), WebApiConfig.ShardingObj.connstring))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
// CHECK SCHEMA
// SQL INJECTION SECURITY ISSUE
cmd.CommandText = "SELECT CompaniesID FROM [mpbdm].[Accounts] JOIN [mpbdm].[Users] ON [mpbdm].[Users].Id = [mpbdm].[Accounts].User_Id WHERE email='" + loginRequest.email + "'";
cmd.CommandType = CommandType.Text;
cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn;
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
// Async
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
bool res = sdr.Read();
if (res != false)
{
shardKey = new Guid(sdr.GetString(0));
}
else
{
return(this.Request.CreateResponse(HttpStatusCode.Unauthorized, "Account doesn't exist!"));
}
}
}
}
// Connect with entity framework to the specific shard
mpbdmContext <Guid> context = new mpbdmContext <Guid>(WebApiConfig.ShardingObj.ShardMap, shardKey, WebApiConfig.ShardingObj.connstring);
Account account = context.Accounts.Include("User").Where(a => a.User.Email == loginRequest.email).SingleOrDefault();
if (account != null)
{
byte[] incoming = CustomLoginProviderUtils.hash(loginRequest.password, account.Salt);
if (CustomLoginProviderUtils.slowEquals(incoming, account.SaltedAndHashedPassword))
{
ClaimsIdentity claimsIdentity = new ClaimsIdentity();
claimsIdentity.AddClaim(new Claim(ClaimTypes.NameIdentifier, account.User.Email));
// Custom Claim must be added to CustomLoginProvider too !!
claimsIdentity.AddClaim(new Claim("shardKey", account.User.CompaniesID));
var customLoginProvider = new CustomLoginProvider(handler);
LoginResult loginResult = customLoginProvider.CreateLoginResult(claimsIdentity, Services.Settings.MasterKey);
MobileLoginResult res = new MobileLoginResult(account, loginResult);
return(this.Request.CreateResponse(HttpStatusCode.OK, res));
}
}
return(this.Request.CreateResponse(HttpStatusCode.Unauthorized, "Invalid username or password"));
}
public void TestShardCommandRetryExhaustion()
{
var retryPolicy = new RetryPolicy(2, TimeSpan.FromMilliseconds(100), TimeSpan.FromSeconds(10), TimeSpan.FromMilliseconds(100));
var shardConnections = new List <Tuple <ShardLocation, DbConnection> >();
// Create ten mocked connections, half of them will throw exceptions on Open
for (int i = 0; i < 10; i++)
{
string database = string.Format("Shard{0}", i);
int j = i;
Action executeOnOpen = () =>
{
if (j < 5)
{
throw new TimeoutException();
}
};
var mockCon = new MockSqlConnection(database, executeOnOpen);
shardConnections.Add(new Tuple <ShardLocation, DbConnection>(new ShardLocation("test", database), mockCon));
}
var mockCmd = new MockSqlCommand();
mockCmd.ExecuteReaderFunc = (t, c) => new MockSqlDataReader();
mockCmd.CommandText = "select 1";
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 300))
{
cmd.ExecutionOptions = MultiShardExecutionOptions.None;
cmd.RetryPolicy = retryPolicy;
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
MultiShardDataReader rdr = cmd.ExecuteReader(CommandBehavior.Default);
// Validate the right exception is re-thrown
Assert.IsTrue(rdr.MultiShardExceptions.Count == 5, "Expected MultiShardExceptions!");
foreach (MultiShardException ex in rdr.MultiShardExceptions)
{
Assert.IsTrue(ex.InnerException is TimeoutException, "Expected TimeoutException!");
}
// Validate that the connections for the faulted readers are closed
for (int i = 0; i < 5; i++)
{
Assert.IsTrue(shardConnections[i].Item2.State == ConnectionState.Closed,
"Expected Connection to be Closed!");
}
}
}
}
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 TestShardCommandCancellation()
{
// Create connections to a few shards
var shardConnections = CreateConnections(10, () => { });
var mockCmd = new MockSqlCommand();
var cmdStartEvent = new ManualResetEvent(false);
mockCmd.ExecuteReaderFunc = (token, cmd) =>
{
while (true)
{
if (token == null)
{
break;
}
token.ThrowIfCancellationRequested();
Thread.Sleep(500);
cmdStartEvent.Set();
}
return(new MockSqlDataReader());
};
mockCmd.CommandText = "select 1";
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 300))
{
try
{
// start the Cancel on a separate thread
Task executeTask = Task.Run(() =>
{
cmdStartEvent.WaitOne();
cmd.Cancel();
});
cmd.ExecuteReader();
executeTask.Wait();
Assert.Fail("We should always be throwing an exception.");
}
catch (Exception ex)
{
Assert.IsTrue(ex is OperationCanceledException, "OperationCanceledException expected. Found {0}!", ex.ToString());
}
}
}
}
public void TestShardCommandRetryBasic()
{
var retryPolicy = new RetryPolicy(4, TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(10), TimeSpan.FromMilliseconds(100));
var openRetryCounts = new int[10];
var shardConnections = new List <Tuple <ShardLocation, DbConnection> >();
// Create ten mocked connections, each will retry retryPolicy.RetryCount - 1 times,
// and keep it's own actual retry count in one of the elements of openRetryCounts
for (int i = 0; i < 10; i++)
{
string database = string.Format("Shard{0}", i);
// We want to close on the value of i
int j = i;
Action executeOnOpen = () =>
{
if (openRetryCounts[j] < (retryPolicy.RetryCount - 1))
{
Logger.Log("Current retry count for database: {0} is {1}", database, openRetryCounts[j]);
openRetryCounts[j]++;
throw new TimeoutException();
}
};
var mockCon = new MockSqlConnection(database, executeOnOpen);
shardConnections.Add(new Tuple <ShardLocation, DbConnection>(new ShardLocation("test", database), mockCon));
}
var mockCmd = new MockSqlCommand();
mockCmd.ExecuteReaderFunc = (t, c) => new MockSqlDataReader();
mockCmd.CommandText = "select 1";
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 300))
{
cmd.ExecutionOptions = MultiShardExecutionOptions.None;
cmd.RetryPolicy = retryPolicy;
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
cmd.ExecuteReader(CommandBehavior.Default);
}
}
for (int i = 0; i < openRetryCounts.Length; i++)
{
Assert.AreEqual(retryPolicy.RetryCount - 1, openRetryCounts[i]);
}
}
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!");
}
}
public void TestSimpleSelect(MultiShardExecutionPolicy policy)
{
// What we're doing:
// Grab all rows from each test database.
// Load them into a MultiShardDataReader.
// Iterate through the rows and make sure that we have 9 total.
//
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = "SELECT dbNameField, Test_int_Field, Test_bigint_Field FROM ConsistentShardedTable";
cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn;
cmd.ExecutionPolicy = policy;
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
int recordsRetrieved = 0;
Logger.Log("Starting to get records");
while (sdr.Read())
{
recordsRetrieved++;
string dbNameField = sdr.GetString(0);
int testIntField = sdr.GetFieldValue <int>(1);
Int64 testBigIntField = sdr.GetFieldValue <Int64>(2);
string shardIdPseudoColumn = sdr.GetFieldValue <string>(3);
string logRecord =
string.Format(
"RecordRetrieved: dbNameField: {0}, TestIntField: {1}, TestBigIntField: {2}, shardIdPseudoColumnField: {3}, RecordCount: {4}",
dbNameField, testIntField, testBigIntField, shardIdPseudoColumn, recordsRetrieved);
Logger.Log(logRecord);
Debug.WriteLine(logRecord);
}
sdr.Close();
Assert.AreEqual(recordsRetrieved, 9);
}
}
}
}
public void TestShardCommandTimeoutException()
{
var shardConnections = CreateConnections(10, () => { });
Func <CancellationToken, MockSqlCommand, DbDataReader> executeReaderFunc = (token, cmd) =>
{
Thread.Sleep(TimeSpan.FromSeconds(2));
return(new MockSqlDataReader());
};
var mockCmd = new MockSqlCommand();
mockCmd.ExecuteReaderFunc = executeReaderFunc;
mockCmd.CommandText = "Select 1";
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 1))
{
cmd.ExecuteReader();
}
}
}
public void TestSelectNonQuery(string commandText, MultiShardExecutionPolicy policy)
{
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = commandText;
cmd.ExecutionPolicy = policy;
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
Assert.AreEqual(0, sdr.MultiShardExceptions.Count);
// TODO: This is a weird error message, but it's good enough for now
// Fixing this will require significant refactoring of MultiShardDataReader,
// we should fix it when we finish implementing async adding of child readers
AssertExtensions.AssertThrows <MultiShardDataReaderClosedException>(() => sdr.Read());
}
}
}
}
public void TestSelect_PartialFailure_PartialResults()
{
using (MultiShardConnection conn = new MultiShardConnection(_shardMap.GetShards(), MultiShardTestUtils.ShardConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = GetPartialFailureQuery();
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
// Exactly one should have failed
Assert.AreEqual(1, sdr.MultiShardExceptions.Count);
// We should be able to read
while (sdr.Read())
{
}
}
}
}
}
public void TestShardCommandFaultHandler()
{
var shardConnections = CreateConnections(10, () => { });
Func <CancellationToken, MockSqlCommand, DbDataReader> executeReaderFunc = (token, cmd) =>
{
throw new InsufficientMemoryException();
};
var mockCmd = new MockSqlCommand();
mockCmd.ExecuteReaderFunc = executeReaderFunc;
mockCmd.CommandText = "Select 1";
ConcurrentDictionary <ShardLocation, bool> passedLocations = new ConcurrentDictionary <ShardLocation, bool>();
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 1))
{
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
cmd.CommandTimeout = 300;
cmd.ShardExecutionFaulted += new EventHandler <ShardExecutionEventArgs>((obj, eventArgs) =>
{
Assert.IsTrue(shardConnections.Select(x => x.Item1).Contains(eventArgs.ShardLocation), "The ShardLocation passed to the event handler does not exist in the set of passed in ShardLocations");
passedLocations[eventArgs.ShardLocation] = true;
Assert.IsInstanceOfType(eventArgs.Exception, typeof(InsufficientMemoryException), "An incorrect exception type was passed to the event handler.");
});
try
{
// We want to execute to completion so we get to the validation at the end of the function.
cmd.ExecuteReader();
}
catch { }
}
}
Assert.AreEqual(shardConnections.Count, passedLocations.Count, "Not every ShardLocation had its corresponding event handler invoked.");
}
public void TestShardResultFailures()
{
ProxyServer proxyServer = GetProxyServer();
try
{
// Start up the proxy server. Do it in a try so we can shut it down in the finally.
// Also, we have to generate the proxyShardconnections *AFTER* we start up the server
// so that we know what port the proxy is listening on. More on the placement
// of the connection generation below.
//
proxyServer.Start();
// PreKillReads is the number of successful reads to perform before killing
// all the connections. We start at 0 to test the no failure case as well.
//
for (int preKillReads = 0; preKillReads <= 10; preKillReads++)
{
// Additionally, since we are running inside a loop, we need to regenerate the proxy shard connections each time
// so that we don't re-use dead connections. If we do that we will end up hung in the read call.
//
List <Tuple <ShardLocation, DbConnection> > proxyShardConnections = GetProxyShardConnections(proxyServer);
using (MultiShardConnection conn = new MultiShardConnection(proxyShardConnections))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = "SELECT db_name() as dbName1, REPLICATE(db_name(), 1000) as longExpr, db_name() as dbName2 FROM ConsistentShardedTable";
cmd.CommandType = CommandType.Text;
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn;
using (MultiShardDataReader sdr = cmd.ExecuteReader(CommandBehavior.Default))
{
int tuplesRead = 0;
while (sdr.Read())
{
// Read part of the tuple first before killing the connections and
// then attempting to read the rest of the tuple.
//
tuplesRead++;
try
{
// The longExpr should contain the first dbName field multiple times.
//
string dbName1 = sdr.GetString(0);
string longExpr = sdr.GetString(1);
Assert.IsTrue(longExpr.Contains(dbName1));
if (tuplesRead == preKillReads)
{
proxyServer.KillAllConnections();
}
// The second dbName field should be the same as the first dbName field.
//
string dbName2 = sdr.GetString(2);
Assert.AreEqual(dbName1, dbName2);
// The shardId should contain both the first and the second dbName fields.
//
string shardId = sdr.GetString(3);
Assert.IsTrue(shardId.Contains(dbName1));
Assert.IsTrue(shardId.Contains(dbName2));
}
catch (Exception ex)
{
// We've seen some failures here due to an attempt to access a socket after it has
// been disposed. The only place where we are attempting to access the socket
// is in the call to proxyServer.KillAllConnections. Unfortunately, it's not clear
// what is causing that problem since it only appears to repro in the lab.
// I (errobins) would rather not blindly start changing things in the code (either
// our code above, our exception handling code here, or the proxyServer code) until
// we know which socket we are trying to access when we hit this problem.
// So, the first step I will take is to pull additional exception information
// so that we can see some more information about what went wrong the next time it repros.
//
Assert.Fail("Unexpected exception, rethrowing. Here is some info: \n Message: {0} \n Source: {1} \n StackTrace: {2}",
ex.Message, ex.Source, ex.StackTrace);
throw;
}
}
Assert.IsTrue((tuplesRead <= preKillReads) || (0 == preKillReads),
String.Format("Tuples read was {0}, Pre-kill reads was {1}", tuplesRead, preKillReads));
}
}
}
}
}
finally
{
// Be sure to shut down the proxy server.
//
string proxyLog = proxyServer.EventLog.ToString();
Logger.Log(proxyLog);
proxyServer.Stop();
}
}
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);
}
public void TestAddDataReaderWithNullSchema()
{
// Creates a MultiShardDataReader and verifies that the right exception is thrown
Func <LabeledDbDataReader[], bool> createMultiShardReader = (readers) =>
{
bool hitNullSchemaException = false;
try
{
var mockMultiShardCmd = MultiShardCommand.Create(null, "test");
var multiShardDataReader = new MultiShardDataReader(mockMultiShardCmd, readers,
MultiShardExecutionPolicy.PartialResults, false, readers.Length);
}
catch (MultiShardDataReaderInternalException ex)
{
hitNullSchemaException = ex.Message.Contains("null schema");
}
return(hitNullSchemaException);
};
var labeledDataReaders = new LabeledDbDataReader[10];
// Create a few mock readers. All with a null schema
for (int i = 0; i < labeledDataReaders.Length; i++)
{
var mockReader = new MockSqlDataReader(string.Format("Reader{0}", i), null /* Null schema */);
labeledDataReaders[i] = new LabeledDbDataReader(mockReader,
new ShardLocation("test", string.Format("Shard{0}", i)), new MockSqlCommand()
{
Connection = new MockSqlConnection("", () => { })
});
}
// Case #1
bool hitException = createMultiShardReader(labeledDataReaders);
Assert.IsFalse(hitException, "Unexpected exception! All readers have a null schema.");
// Case #2
for (int i = 0; i < labeledDataReaders.Length; i++)
{
MockSqlDataReader mockReader = (MockSqlDataReader)labeledDataReaders[i].DbDataReader;
mockReader.Open();
if (i > labeledDataReaders.Length / 2)
{
mockReader.DataTable = new DataTable();
}
}
hitException = createMultiShardReader(labeledDataReaders);
Assert.IsTrue(hitException, "Exception not hit! Second half of readers don't have a null schema!");
// Case #3
for (int i = 0; i < labeledDataReaders.Length; i++)
{
MockSqlDataReader mockReader = (MockSqlDataReader)labeledDataReaders[i].DbDataReader;
mockReader.Open();
if (i < labeledDataReaders.Length / 2)
{
mockReader.DataTable = new DataTable();
}
else
{
mockReader.DataTable = null;
}
}
hitException = createMultiShardReader(labeledDataReaders);
Assert.IsTrue(hitException, "Exception not hit! First half of readers don't have a null schema!");
}
public void TestDataReaderReadException()
{
// Setup two data readers from shards
var mockReader1 = new MockSqlDataReader("Reader1");
var mockReader2 = new MockSqlDataReader("Reader2");
bool movedOnToNextReader = false;
int invokeCount = 1;
Func <MockSqlDataReader, Task <bool> > ExecuteOnReadAsync = (r) =>
{
return(Task.Run <bool>(() =>
{
// First reader throws an exception when Read
if (r.Name == "Reader1")
{
if (invokeCount == 2)
{
throw new InvalidOperationException();
}
}
else
{
movedOnToNextReader = true;
}
return true;
}));
};
Action <MockSqlDataReader> ExecuteOnGetColumn = (r) =>
{
if (r.Name == "Reader1")
{
throw new InvalidOperationException();
}
};
mockReader1.ExecuteOnReadAsync = ExecuteOnReadAsync;
mockReader1.ExecuteOnGetColumn = ExecuteOnGetColumn;
mockReader2.ExecuteOnReadAsync = ExecuteOnReadAsync;
var labeledDataReaders = new LabeledDbDataReader[2];
labeledDataReaders[0] = new LabeledDbDataReader(mockReader1, new ShardLocation("test", "Shard1"),
new MockSqlCommand()
{
Connection = new MockSqlConnection("", () => { })
});
labeledDataReaders[1] = new LabeledDbDataReader(mockReader2, new ShardLocation("test", "Shard2"),
new MockSqlCommand()
{
Connection = new MockSqlConnection("", () => { })
});
// Create the MultiShardDataReader
var mockMultiShardCmd = MultiShardCommand.Create(null, "test");
var multiShardDataReader = new MultiShardDataReader(mockMultiShardCmd, labeledDataReaders,
MultiShardExecutionPolicy.PartialResults, false);
// Validate that if an exception is thrown when reading a column,
// it is propagated back to the user
try
{
multiShardDataReader.Read();
invokeCount++;
multiShardDataReader.GetInt32(0);
}
catch (Exception ex)
{
Assert.IsTrue(ex is InvalidOperationException, "Expected InvalidOperationException!");
}
// Validate that we didn't automatically move on to the next reader when we
// hit an exception whilst reading the column and that
// an exception from a second Read() call is stored and the reader is closed
multiShardDataReader.Read();
Assert.AreEqual(multiShardDataReader.MultiShardExceptions.Count, 1, "Expected exception to be recorded");
Assert.IsTrue(mockReader1.IsClosed, "Expected reader to be closed!");
// Validate we immediately moved on to the next reader
multiShardDataReader.Read();
Assert.IsTrue(movedOnToNextReader, "Should've moved on to next reader");
}
public void TestShardCommandRetryConnectionReopen()
{
var retryPolicy = new RetryPolicy(4, TimeSpan.FromMilliseconds(100), TimeSpan.FromSeconds(10), TimeSpan.FromMilliseconds(100));
var shardConnections = new List <Tuple <ShardLocation, DbConnection> >();
// Callback to execute when the MockCommand is invoked
Func <CancellationToken, MockSqlCommand, DbDataReader> ExecuteReaderFunc = null;
// Number of times each command has been retried
var commandRetryCounts = new int[10];
// Create ten mocked connections,
// a few of them will throw an exception on Open
// and the rest will throw an exception on command execution upto 2 retries
// At the end, all commands should complete successfully.
for (int i = 0; i < 10; i++)
{
string database = string.Format("{0}", i);
int j = i;
int retryCount = 0;
Action executeOnOpen = () =>
{
if (j < 5)
{
if (retryCount < 3)
{
retryCount++;
throw new TimeoutException();
}
}
};
var mockCon = new MockSqlConnection(database, executeOnOpen);
shardConnections.Add(new Tuple <ShardLocation, DbConnection>(new ShardLocation("Shard", database), mockCon));
}
ExecuteReaderFunc = (t, r) =>
{
int index = Int32.Parse(((MockSqlConnection)(r.Connection)).ConnectionString);
if (r.Connection.State == ConnectionState.Closed)
{
throw new InvalidOperationException("Command shouldn't be executed on a closed connection!");
}
if (index > 5 && commandRetryCounts[index] < 3)
{
commandRetryCounts[index]++;
r.RetryCount++;
r.Connection.Close();
throw new TimeoutException();
}
else
{
var mockRdr = new MockSqlDataReader();
mockRdr.ExecuteOnReadAsync = (rdr) =>
{
return(Task.Run <bool>(() =>
{
bool isClosed = rdr.IsClosed;
rdr.Close();
return !isClosed;
}));
};
return(mockRdr);
}
};
var mockCmd = new MockSqlCommand();
mockCmd.ExecuteReaderFunc = ExecuteReaderFunc;
mockCmd.CommandText = "select 1";
using (var conn = new MultiShardConnection(shardConnections))
{
using (var cmd = MultiShardCommand.Create(conn, mockCmd, 300))
{
cmd.RetryPolicy = retryPolicy;
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
using (var reader = cmd.ExecuteReaderAsync().Result)
{
// Validate that we successfully received a reader
// from each one of the shards
int readerCount = 0;
while (reader.Read())
{
readerCount++;
}
Assert.AreEqual(10, readerCount, "Expected 10 readers!");
}
}
}
}
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));
}
/// <summary>
/// Creates and returns a <see cref="MultiShardCommand"/> object.
/// The <see cref="MultiShardCommand"/> object can then be used to
/// execute a command against all shards specified in the connection.
/// </summary>
/// <returns>the <see cref="MultiShardCommand"/> with <see cref="MultiShardCommand.CommandText"/> set to null.</returns>
public MultiShardCommand CreateCommand()
{
return(MultiShardCommand.Create(this, commandText: null));
}
public List <string[]> ExecuteMultiShardQuery(ListShardMap <int> shardMap, string credentialsConnectionString)
{
// Get the shards to connect to
IEnumerable <Shard> shards = shardMap.GetShards();
// Create the multi-shard connection
using (MultiShardConnection conn = new MultiShardConnection(shards, credentialsConnectionString))
{
// Create a simple command
using (MultiShardCommand cmd = conn.CreateCommand())
{
// Because this query is grouped by CustomerID, which is sharded,
// we will not get duplicate rows.
cmd.CommandText = @"
SELECT
c.CustomerId,
c.Name AS CustomerName,
o.OrderID AS OrderCount
FROM
dbo.Customers AS c INNER JOIN
dbo.Orders AS o
ON c.CustomerID = o.CustomerID
ORDER BY
1";
// Append a column with the shard name where the row came from
cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn;
// Allow for partial results in case some shards do not respond in time
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
// Allow the entire command to take up to 30 seconds
cmd.CommandTimeout = 30;
// Execute the command.
// We do not need to specify retry logic because MultiShardDataReader will internally retry until the CommandTimeout expires.
using (MultiShardDataReader reader = cmd.ExecuteReader())
{
// Get the column names
TableFormatter formatter = new TableFormatter(GetColumnNames(reader).ToArray());
int rows = 0;
while (reader.Read())
{
// Read the values using standard DbDataReader methods
object[] values = new object[reader.FieldCount];
reader.GetValues(values);
// Extract just database name from the $ShardLocation pseudocolumn to make the output formater cleaner.
// Note that the $ShardLocation pseudocolumn is always the last column
int shardLocationOrdinal = values.Length - 1;
values[shardLocationOrdinal] = ExtractDatabaseName(values[shardLocationOrdinal].ToString());
// Add values to output formatter
formatter.AddRow(values);
rows++;
}
Console.WriteLine(formatter.ToString());
Console.WriteLine("({0} rows returned)", rows);
return(formatter._rows);
}
}
}
}
public static void Main()
{
SqlConnectionStringBuilder connStrBldr = new SqlConnectionStringBuilder
{
UserID = userName,
Password = password,
ApplicationName = applicationName
};
#region Create Shards
// Bootstrap the shard map manager, register shards, and store mappings of tenants to shards
// Note that you can keep working with existing shard maps. There is no need to
// re-create and populate the shard map from scratch every time.
Console.WriteLine("Checking for existing shard map and creating new shard map if necessary.");
Sharding sharding = new Sharding(server, shardmapmgrdb, connStrBldr.ConnectionString);
sharding.RegisterNewShard(server, shard1, connStrBldr.ConnectionString, tenantId1);
sharding.RegisterNewShard(server, shard2, connStrBldr.ConnectionString, tenantId2);
#endregion
#region Insert Items
// Do work for tenant 1 :-)
// Create and save a new Blog
Console.Write("Enter a name for a new Blog: ");
var name = Console.ReadLine();
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (var db = new ElasticScaleContext <int>(sharding.ShardMap, tenantId1, connStrBldr.ConnectionString))
{
var blog = new Blog {
Name = name
};
db.Blogs.Add(blog);
db.SaveChanges();
}
});
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (var db = new ElasticScaleContext <int>(sharding.ShardMap, tenantId1, connStrBldr.ConnectionString))
{
// Display all Blogs for tenant 1
var query = from b in db.Blogs
orderby b.Name
select b;
Console.WriteLine("All blogs for tenant id {0}:", tenantId1);
foreach (var item in query)
{
Console.WriteLine(item.Name);
}
}
});
// Do work for tenant 2 :-)
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (var db = new ElasticScaleContext <int>(sharding.ShardMap, tenantId2, connStrBldr.ConnectionString))
{
// Display all Blogs from the database
var query = from b in db.Blogs
orderby b.Name
select b;
Console.WriteLine("All blogs for tenant id {0}:", tenantId2);
foreach (var item in query)
{
Console.WriteLine(item.Name);
}
}
});
// Create and save a new Blog
Console.Write("Enter a name for a new Blog: ");
var name2 = Console.ReadLine();
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (var db = new ElasticScaleContext <int>(sharding.ShardMap, tenantId2, connStrBldr.ConnectionString))
{
var blog = new Blog {
Name = name2
};
db.Blogs.Add(blog);
db.SaveChanges();
}
});
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (var db = new ElasticScaleContext <int>(sharding.ShardMap, tenantId2, connStrBldr.ConnectionString))
{
// Display all Blogs from the database
var query = from b in db.Blogs
orderby b.Name
select b;
Console.WriteLine("All blogs for tenant id {0}:", tenantId2);
foreach (var item in query)
{
Console.WriteLine(item.Name);
}
}
});
#endregion
#region Query
Console.WriteLine("Staring Multi Shard Read via MultiShardCommand");
/// Multi Shard querying
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (MultiShardConnection conn = new MultiShardConnection(
sharding.ShardMap.GetShards(),
connStrBldr.ConnectionString))
{
using (MultiShardCommand cmd = conn.CreateCommand())
{
cmd.CommandText = "SELECT BlogId, Name FROM Blogs";
cmd.CommandType = CommandType.Text;
cmd.ExecutionOptions = MultiShardExecutionOptions.IncludeShardNameColumn;
cmd.ExecutionPolicy = MultiShardExecutionPolicy.PartialResults;
using (MultiShardDataReader sdr = cmd.ExecuteReader())
{
while (sdr.Read())
{
var c2Field = sdr.GetFieldValue <int>(0);
var c1Field = sdr.GetFieldValue <string>(1);
Blog blog = new Blog(c2Field, c1Field, null);
Console.WriteLine("Blog Entry: " + blog.Name);
}
}
}
}
});
Console.WriteLine("Staring Multi Shard Read via Elastic Scale");
//Elastic Query
SqlDatabaseUtils.SqlRetryPolicy.ExecuteAction(() =>
{
using (SqlConnection conn = new SqlConnection(elasticqueryString))
{
using (SqlCommand cmd = conn.CreateCommand())
{
cmd.CommandText = "SELECT BlogId, Name FROM Blogs";
cmd.CommandType = CommandType.Text;
conn.Open();
using (SqlDataReader sdr = cmd.ExecuteReader())
{
while (sdr.Read())
{
var c2Field = sdr.GetFieldValue <int>(0);
var c1Field = sdr.GetFieldValue <string>(1);
Blog blog = new Blog(c2Field, c1Field, null);
Console.WriteLine("Blog Entry: " + blog.Name);
}
}
}
}
});
#endregion
Console.WriteLine("Press any key to exit...");
Console.ReadKey();
}
