Expose MultiShardConnection with AccessToken support (#201)

* Refresh ShardConnections access tokens
* Expose MultiShardConnection with SqlConnectionInfo

Author: jaredmoo

Src/ElasticScale.Client/Query/MultiShardCommand.cs

@@ -447,26 +447,12 @@ internal MultiShardDataReader ExecuteReader(
         {
             Contract.Requires(commandRetryPolicy != null && connectionRetryPolicy != null);
 
-            try
-            {
-                return this.ExecuteReaderAsync(
-                    behavior,
-                    CancellationToken.None,
-                    commandRetryPolicy,
-                    connectionRetryPolicy,
-                    executionPolicy).Result;
-            }
-            catch (Exception ex)
-            {
-                AggregateException aex = ex as AggregateException;
-
-                if (null != aex)
-                {
-                    throw aex.Flatten().InnerException;
-                }
-
-                throw;
-            }
+            return this.ExecuteReaderAsync(
+                behavior,
+                CancellationToken.None,
+                commandRetryPolicy,
+                connectionRetryPolicy,
+                executionPolicy).GetAwaiter().GetResult();
         }
 
 #endregion
@@ -1235,7 +1221,7 @@ private bool IsExecutionInProgress()
         private List<Tuple<ShardLocation, DbCommand>> GetShardDbCommands()
         {
             return this.Connection
-                       .ShardConnections
+                       .GetShardConnections()
                        .Select(sc => new Tuple<ShardLocation, DbCommand>(sc.Item1, MultiShardUtils.CloneDbCommand(_dbCommand, sc.Item2)))
                        .ToList();
         }

Src/ElasticScale.Client/Query/MultiShardConnection.cs

@@ -49,7 +49,12 @@ public sealed class MultiShardConnection : IDisposable
         /// </summary>
         private SqlConnectionInfo _connectionInfo;
 
-#endregion
+        /// <summary>
+        /// The shard connections
+        /// </summary>
+        private List<Tuple<ShardLocation, DbConnection>> _shardConnections;
+
+        #endregion
 
 #region Ctors
 
@@ -89,14 +94,50 @@ public MultiShardConnection(IEnumerable<ShardLocation> shardLocations, string co
             InitializeShardConnections(shardLocations);
         }
 
+        /// <summary>
+        /// Initializes a new instance of the <see cref="MultiShardConnection"/> class.
+        /// </summary>
+        /// <param name="shards">The collection of <see cref="Shard"/>s used for this connection instances.</param>
+        /// <param name="connectionInfo">
+        /// These credentials will be used to connect to the <see cref="Shard"/>s.
+        /// The same credentials are used on all shards.
+        /// Therefore, all shards need to provide the appropriate permissions for these credentials to execute the command.
+        /// </param>
+        /// <remarks>
+        /// Multiple Active Result Sets (MARS) are not supported and are disabled for any processing at the shards.
+        /// </remarks>
+        public MultiShardConnection(IEnumerable<Shard> shards, SqlConnectionInfo connectionInfo)
+        {
+            InitializeConnectionInfo(connectionInfo);
+            InitializeShardConnections(shards);
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="MultiShardConnection"/> class.
+        /// </summary>
+        /// <param name="shardLocations">The collection of <see cref="ShardLocation"/>s used for this connection instances.</param>
+        /// <param name="connectionInfo">
+        /// These credentials will be used to connect to the <see cref="Shard"/>s.
+        /// The same credentials are used on all shards.
+        /// Therefore, all shards need to provide the appropriate permissions for these credentials to execute the command.
+        /// </param>
+        /// <remarks>
+        /// Multiple Active Result Sets (MARS) are not supported and are disabled for any processing at the shards.
+        /// </remarks>
+        public MultiShardConnection(IEnumerable<ShardLocation> shardLocations, SqlConnectionInfo connectionInfo)
+        {
+            InitializeConnectionInfo(connectionInfo);
+            InitializeShardConnections(shardLocations);
+        }
+
         /// <summary>
         /// Creates an instance of this class
         /// /* TEST ONLY */
         /// </summary>
         /// <param name="shardConnections">Connections to the shards</param>
         internal MultiShardConnection(List<Tuple<ShardLocation, DbConnection>> shardConnections)
         {
-            this.ShardConnections = shardConnections;
+            this._shardConnections = shardConnections;
         }
 
 #endregion
@@ -119,17 +160,11 @@ public IEnumerable<ShardLocation> ShardLocations
         {
             get
             {
-                return this.ShardConnections.Select(s => s.Item1);
+                return this._shardConnections.Select(s => s.Item1);
             }
         }
 
-        internal List<Tuple<ShardLocation, DbConnection>> ShardConnections
-        {
-            get;
-            private set;
-        }
-
-#endregion
+        #endregion
 
 #region Public Methods
 
@@ -152,7 +187,7 @@ public void Dispose()
             if (!_disposed)
             {
                 // Dispose off the shard connections
-                this.ShardConnections.ForEach(
+                this._shardConnections.ForEach(
                 (c) =>
                 {
                     if (c.Item2 != null)
@@ -167,9 +202,35 @@ public void Dispose()
             }
         }
 
-#endregion
+        #endregion
+
+
+        #region Internal methods
+
+        /// <summary>
+        /// Gets the shard connections
+        /// </summary>
+        internal List<Tuple<ShardLocation, DbConnection>> GetShardConnections()
+        {
+            // Refresh the access tokens for all shard connections.
+            // (Null check because unit tests use internal code path which doesn't initialize _connectionInfo).
+            if (this._connectionInfo != null)
+            {
+                foreach (var shardConnection in _shardConnections)
+                {
+                    if (shardConnection.Item2 != null)
+                    {
+                        this._connectionInfo.RefreshAccessToken(shardConnection.Item2);
+                    }
+                }
+            }
 
-#region Helpers
+            return _shardConnections;
+        }
+
+        #endregion
+
+        #region Helpers
 
         private static void ValidateNotEmpty<T>(
             IEnumerable<T> namedCollection,
@@ -181,11 +242,28 @@ private static void ValidateNotEmpty<T>(
             }
         }
 
+        private void InitializeConnectionInfo(SqlConnectionInfo connectionInfo)
+        {
+            if (connectionInfo == null)
+            {
+                throw new ArgumentNullException("connectionInfo");
+            }
+
+            string updatedConnectionString = InitializeConnectionString(connectionInfo.ConnectionString);
+            this._connectionInfo = connectionInfo.CloneWithUpdatedConnectionString(updatedConnectionString);
+        }
+
         private void InitializeConnectionInfo(string connectionString)
+        {
+            string updatedConnectionString = InitializeConnectionString(connectionString);
+            this._connectionInfo = new SqlConnectionInfo(updatedConnectionString);
+        }
+
+        private static string InitializeConnectionString(string connectionString)
         {
             if (connectionString == null)
             {
-                throw new ArgumentNullException("connectionString");
+                throw new ArgumentNullException(nameof(connectionString));
             }
 
             // Enhance the ApplicationName with this library's name as a suffix
@@ -206,7 +284,7 @@ private void InitializeConnectionInfo(string connectionString)
                 throw new ArgumentException("InitialCatalog must not be set in the connectionStringBuilder");
             }
 
-            this._connectionInfo = new SqlConnectionInfo(connectionStringBuilder.ToString());
+            return connectionStringBuilder.ToString();
         }
 
         private void InitializeShardConnections(IEnumerable<Shard> shards)
@@ -220,8 +298,8 @@ private void InitializeShardConnections(IEnumerable<Shard> shards)
             this.Shards = shards.ToList();
             ValidateNotEmpty(this.Shards, "shards");
 
-            this.ShardConnections = this.Shards.Select(
-                s => CreateDbConnectionForLocation(s.Location, _connectionInfo)).ToList();
+            this._shardConnections = (this.Shards.Select(
+                s => CreateDbConnectionForLocation(s.Location, _connectionInfo)).ToList());
         }
 
         private void InitializeShardConnections(IEnumerable<ShardLocation> shardLocations)
@@ -236,7 +314,7 @@ private void InitializeShardConnections(IEnumerable<ShardLocation> shardLocation
             ValidateNotEmpty(shardLocationsList, "shardLocations");
 
             this.Shards = null;
-            this.ShardConnections = shardLocationsList.Select(
+            this._shardConnections = shardLocationsList.Select(
                 s => CreateDbConnectionForLocation(s, _connectionInfo)).ToList();
         }
 
@@ -270,7 +348,7 @@ private static Tuple<ShardLocation, DbConnection> CreateDbConnectionForLocation(
         [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "We do not want to throw on Close.")]
         internal void Close()
         {
-            foreach (var conn in this.ShardConnections)
+            foreach (var conn in this._shardConnections)
             {
                 if (conn.Item2 != null && conn.Item2.State != ConnectionState.Closed)
                 {

Src/ElasticScale.Client/ShardManagement/SqlStore/SqlConnectionInfo.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Data.Common;
 using System.Data.SqlClient;
 
 namespace Microsoft.Azure.SqlDatabase.ElasticScale.ShardManagement
@@ -90,6 +91,22 @@ internal SqlConnection CreateConnection()
                 Credential = Credential,
             };
 
+            RefreshAccessToken(conn);
+
+            return conn;
+        }
+
+        internal void RefreshAccessToken(DbConnection conn)
+        {
+            SqlConnection sqlConn= conn as SqlConnection;
+            if (sqlConn != null)
+            {
+                RefreshAccessToken(sqlConn);
+            }
+        }
+
+        internal void RefreshAccessToken(SqlConnection conn)
+        {
             if (AccessTokenFactory != null)
             {
 #if NET451
@@ -98,8 +115,6 @@ internal SqlConnection CreateConnection()
                 conn.AccessToken = AccessTokenFactory();
 #endif
             }
-
-            return conn;
         }
 
         /// <summary>

Test/ElasticScale.Query.UnitTests/MultiShardDataReaderTests.cs

@@ -125,15 +125,15 @@ public void MyTestInitialize()
             //          BUT, since we are writing tests at a lower level than that, we need to open
             //          the connections manually here.  Hence the loop below.
             //
-            foreach (var conn in _shardConnection.ShardConnections)
+            foreach (var conn in _shardConnection.GetShardConnections())
             {
                 conn.Item2.Open();
             }
 
-            _conn1 = (SqlConnection)_shardConnection.ShardConnections[0].Item2;
-            _conn2 = (SqlConnection)_shardConnection.ShardConnections[1].Item2;
-            _conn3 = (SqlConnection)_shardConnection.ShardConnections[2].Item2;
-            _conns = _shardConnection.ShardConnections.Select(x => (SqlConnection)x.Item2);
+            _conn1 = (SqlConnection)_shardConnection.GetShardConnections()[0].Item2;
+            _conn2 = (SqlConnection)_shardConnection.GetShardConnections()[1].Item2;
+            _conn3 = (SqlConnection)_shardConnection.GetShardConnections()[2].Item2;
+            _conns = _shardConnection.GetShardConnections().Select(x => (SqlConnection)x.Item2);
         }
 
         /// <summary>
@@ -142,7 +142,7 @@ public void MyTestInitialize()
         [TestCleanup()]
         public void MyTestCleanup()
         {
-            foreach (var conn in _shardConnection.ShardConnections)
+            foreach (var conn in _shardConnection.GetShardConnections())
             {
                 conn.Item2.Close();
             }