Query.Store() works in async queries

Author: alex-kulakov

Orm/Xtensive.Orm.Tests/Linq/QueryMethodTests.cs

@@ -6,6 +6,7 @@
 
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using NUnit.Framework;
 using Xtensive.Collections;
 using Xtensive.Core;
@@ -119,6 +120,29 @@ public void Store1Test()
       Assert.AreEqual(0, expected.Except(query).Count());
     }
 
+    [Test]
+    public async Task Store1TestAsync()
+    {
+      Require.AllFeaturesSupported(ProviderFeatures.TemporaryTables);
+
+      var localCustomers = Session.Query.All<Customer>().Take(10).ToList();
+      var query = (await Session.Query.All<Customer>()
+        .Join(
+          Session.Query.Store(localCustomers),
+          customer => customer,
+          localCustomer => localCustomer,
+          (customer, localCustomer) => new { customer, localCustomer }).AsAsync()).ToList();
+      var expected = Session.Query.All<Customer>().AsEnumerable()
+        .Join(
+          Session.Query.Store(localCustomers),
+          customer => customer,
+          localCustomer => localCustomer,
+          (customer, localCustomer) => new { customer, localCustomer });
+
+      Assert.That(query, Is.Not.Empty);
+      Assert.AreEqual(0, expected.Except(query).Count());
+    }
+
     [Test]
     public void Store2Test()
     {

Orm/Xtensive.Orm/Orm/Providers/Interfaces/IProviderExecutor.cs

@@ -38,6 +38,13 @@ public interface IProviderExecutor
     /// <param name="tuples">The tuples to store.</param>
     void Store(IPersistDescriptor descriptor, IEnumerable<Tuple> tuples);
 
+    /// <summary>
+    /// Asynchronously stores the specified tuples in specified table.
+    /// </summary>
+    /// <param name="descriptor">Persist descriptor.</param>
+    /// <param name="tuples">The tuples to store.</param>
+    Task StoreAsync(EnumerationContext context, IPersistDescriptor descriptor, IEnumerable<Tuple> tuples, CancellationToken token);
+
     /// <summary>
     /// Clears the specified table.
     /// </summary>

Orm/Xtensive.Orm/Orm/Providers/SqlSessionHandler.IProviderExecutor.cs

@@ -7,6 +7,7 @@
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using Xtensive.Orm.Rse.Providers;
 using Tuple = Xtensive.Tuples.Tuple;
 
 namespace Xtensive.Orm.Providers
@@ -45,6 +46,27 @@ void IProviderExecutor.Store(IPersistDescriptor descriptor, IEnumerable<Tuple> t
         commandProcessor.ExecuteTasks(context);
     }
 
+
+
+    async Task IProviderExecutor.StoreAsync(EnumerationContext enumerationContext,IPersistDescriptor descriptor, IEnumerable<Tuple> tuples, CancellationToken token)
+    {
+      await PrepareAsync(token);
+
+      if (tuples is ExecutableRawProvider rawProvider) {
+        var enumerator = await rawProvider.GetEnumeratorAsync(enumerationContext, token);
+        while(enumerator.MoveNext()) {
+          commandProcessor.RegisterTask(new SqlPersistTask(descriptor.StoreRequest, enumerator.Current));
+        }
+      }
+      else {
+        foreach (var tuple in tuples)
+          commandProcessor.RegisterTask(new SqlPersistTask(descriptor.StoreRequest, tuple));
+      }
+
+      using (var context = new CommandProcessorContext())
+        await commandProcessor.ExecuteTasksAsync(context, token);
+    }
+
     /// <inheritdoc/>
     void IProviderExecutor.Clear(IPersistDescriptor descriptor)
     {

Orm/Xtensive.Orm/Orm/Providers/SqlSessionHandler.cs

@@ -6,6 +6,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Data;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -162,7 +163,8 @@ private void Prepare()
     private async Task PrepareAsync(CancellationToken cancellationToken)
     {
       Session.EnsureNotDisposed();
-      await driver.OpenConnectionAsync(Session, connection, cancellationToken).ConfigureAwait(false);
+      if (connection.State != ConnectionState.Open)
+        await driver.OpenConnectionAsync(Session, connection, cancellationToken).ConfigureAwait(false);
 
       try {
         foreach (var initializationSqlScript in initializationSqlScripts)

Orm/Xtensive.Orm/Orm/Providers/SqlStoreProvider.cs

@@ -4,6 +4,8 @@
 // Created by: Dmitri Maximov
 // Created:    2008.09.05
 
+using System.Threading;
+using System.Threading.Tasks;
 using Xtensive.Orm.Rse.Providers;
 
 namespace Xtensive.Orm.Providers
@@ -30,6 +32,12 @@ protected override void OnBeforeEnumerate(Rse.Providers.EnumerationContext conte
       LockAndStore(context, Source);
     }
 
+    protected override async Task OnBeforeEnumerateAsync(Rse.Providers.EnumerationContext context, CancellationToken token)
+    {
+      await base.OnBeforeEnumerateAsync(context, token);
+      await LockAndStoreAsync(context, Source, token);
+    }
+
     protected override void OnAfterEnumerate(Rse.Providers.EnumerationContext context)
     {
       ClearAndUnlock(context);

Orm/Xtensive.Orm/Orm/Providers/SqlTemporaryDataProvider.cs

@@ -10,6 +10,8 @@
 using Xtensive.Tuples;
 using Tuple = Xtensive.Tuples.Tuple;
 using Xtensive.Orm.Rse.Providers;
+using System.Threading.Tasks;
+using System.Threading;
 
 namespace Xtensive.Orm.Providers
 {
@@ -33,6 +35,17 @@ protected void LockAndStore(Rse.Providers.EnumerationContext context, IEnumerabl
       executor.Store(tableDescriptor, data);
     }
 
+    protected async Task LockAndStoreAsync(Rse.Providers.EnumerationContext context, IEnumerable<Tuple> data, CancellationToken token)
+    {
+      var storageContext = (EnumerationContext) context;
+      var tableLock = DomainHandler.TemporaryTableManager.Acquire(storageContext, tableDescriptor);
+      if (tableLock == null)
+        return;
+      storageContext.SetValue(this, TemporaryTableLockName, tableLock);
+      var executor = storageContext.Session.Services.Demand<IProviderExecutor>();
+      await executor.StoreAsync(storageContext, tableDescriptor, data, token);
+    }
+
     protected bool ClearAndUnlock(Rse.Providers.EnumerationContext context)
     {
       var tableLock = context.GetValue<IDisposable>(this, TemporaryTableLockName);

Orm/Xtensive.Orm/Orm/Rse/Providers/Executable/ExecutableRawProvider.cs

@@ -6,6 +6,8 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Tuple = Xtensive.Tuples.Tuple;
 
 namespace Xtensive.Orm.Rse.Providers
@@ -16,28 +18,26 @@ namespace Xtensive.Orm.Rse.Providers
   [Serializable]
   public sealed class ExecutableRawProvider : ExecutableProvider<Providers.RawProvider>
   {
-    #region Cached properties
-
     private const string CachedSourceName = "CachedSource";
 
-    private IEnumerable<Tuple> CachedSource {
-      get { return GetValue<IEnumerable<Tuple>>(EnumerationContext.Current, CachedSourceName); }
-      set { SetValue(EnumerationContext.Current, CachedSourceName, value); }
-    }
-
-    #endregion
-
     /// <inheritdoc/>
     protected override void OnBeforeEnumerate(EnumerationContext context)
     {
       base.OnBeforeEnumerate(context);
-      CachedSource = Origin.CompiledSource.Invoke();
+      context.SetValue(this, CachedSourceName, Origin.CompiledSource.Invoke());
+    }
+
+    /// <inheritdoc/>
+    protected override async Task OnBeforeEnumerateAsync(EnumerationContext context, CancellationToken token)
+    {
+      await base.OnBeforeEnumerateAsync(context, token);
+      context.SetValue(this, CachedSourceName, Origin.CompiledSource.Invoke());
     }
 
     /// <inheritdoc/>
     protected override IEnumerable<Tuple> OnEnumerate(EnumerationContext context)
     {
-      return CachedSource;
+      return context.GetValue<IEnumerable<Tuple>>(this, CachedSourceName);
     }
 
 

Orm/Xtensive.Orm/Orm/Rse/Providers/ExecutableProvider.cs

@@ -41,6 +41,16 @@ protected virtual void OnBeforeEnumerate(EnumerationContext context)
       }
     }
 
+    protected virtual async Task OnBeforeEnumerateAsync(EnumerationContext context, CancellationToken token)
+    {
+      token.ThrowIfCancellationRequested();
+      foreach (var source in Sources) {
+        var ep = source as ExecutableProvider;
+        if (ep != null)
+          await ep.OnBeforeEnumerateAsync(context, token);
+      }
+    }
+
     /// <summary>
     /// Called when enumeration is finished.
     /// </summary>
@@ -117,7 +127,7 @@ public async Task<IEnumerator<Tuple>> GetEnumeratorAsync(EnumerationContext cont
       var enumerated = context.GetValue<bool>(this, enumerationMarker);
       bool onEnumerationExecuted = false;
       if (!enumerated)
-        OnBeforeEnumerate(context);
+        await OnBeforeEnumerateAsync(context, token);
       try {
         context.SetValue(this, enumerationMarker, true);
         var enumerator = (await OnEnumerateAsync(context, token).ConfigureAwait(false))