Add AddOrReplaceComponent and use in From

GetOneComponent prioritizes engine

Author: asherber

QueryBuilder.Tests/GeneralTests.cs

@@ -1,6 +1,8 @@
 using SqlKata.Compilers;
 using SqlKata.Extensions;
 using SqlKata.Tests.Infrastructure;
+using System;
+using System.Linq;
 using Xunit;
 
 namespace SqlKata.Tests
@@ -205,5 +207,64 @@ public void CompilerSpecificFromMixed()
             Assert.Equal("SELECT * FROM \"pgsql\"", c[EngineCodes.PostgreSql].RawSql);
             Assert.Equal("SELECT * FROM `mysql`", c[EngineCodes.MySql].RawSql);
         }
+
+        [Fact]
+        public void OneFromPerEngine()
+        {
+            var query = new Query("generic")
+                .ForSqlServer(q => q.From("dnu"))
+                .ForSqlServer(q => q.From("mssql"));
+            var engines = new[] { EngineCodes.SqlServer, EngineCodes.MySql, EngineCodes.PostgreSql };
+            var c = Compilers.Compile(engines, query);
+
+            Assert.Equal(2, query.Clauses.OfType<AbstractFrom>().Count());
+            Assert.Equal("SELECT * FROM [mssql]", c[EngineCodes.SqlServer].RawSql);
+            Assert.Equal("SELECT * FROM \"generic\"", c[EngineCodes.PostgreSql].RawSql);
+            Assert.Equal("SELECT * FROM `generic`", c[EngineCodes.MySql].RawSql);
+        }
+
+        [Theory]
+        [InlineData(null, null)]
+        [InlineData(null, "mssql")]
+        [InlineData("original", null)]
+        [InlineData("original", "mssql")]
+        public void AddOrReplace_Works(string table, string engine)
+        {
+            var query = new Query();
+            if (table != null)
+                query.From(table);
+            query.AddOrReplaceComponent("from", new FromClause() { Table = "updated", Engine = engine });
+            var froms = query.Clauses.OfType<FromClause>();
+
+            Assert.Single(froms);
+            Assert.Equal("updated", froms.Single().Table);
+        }
+
+        [Theory]
+        [InlineData(null, "generic")]
+        [InlineData(EngineCodes.SqlServer, "mssql")]
+        [InlineData(EngineCodes.MySql, "generic")]
+        public void GetOneComponent_Prefers_Engine(string engine, string column)
+        {
+            var query = new Query()
+                .Where("generic", "foo")
+                .ForSqlServer(q => q.Where("mssql", "foo"));
+
+            var where = query.GetOneComponent("where", engine) as BasicCondition;
+
+            Assert.NotNull(where);
+            Assert.Equal(column, where.Column);
+        }
+
+        [Fact]
+        public void AddOrReplace_Throws_MoreThanOne()
+        {
+            var query = new Query()
+                .Where("a", "b")
+                .Where("c", "d");
+
+            Action act = () => query.AddOrReplaceComponent("where", new BasicCondition());
+            Assert.Throws<InvalidOperationException>(act);
+        }
     }
 }

QueryBuilder/BaseQuery.cs

@@ -82,6 +82,28 @@ public Q AddComponent(string component, AbstractClause clause, string engineCode
             return (Q)this;
         }
 
+        /// <summary>
+        /// If the query already contains a clause for the given component
+        /// and engine, replace it with the specified clause. Otherwise, just 
+        /// add the clause.
+        /// </summary>
+        /// <param name="component"></param>
+        /// <param name="clause"></param>
+        /// <param name="engineCode"></param>
+        /// <returns></returns>
+        public Q AddOrReplaceComponent(string component, AbstractClause clause, string engineCode = null)
+        {
+            engineCode = engineCode ?? EngineScope;
+
+            var current = GetComponents(component).SingleOrDefault(c => c.Engine == engineCode);
+            if (current != null)
+                Clauses.Remove(current);
+
+            return AddComponent(component, clause, engineCode);            
+        }
+
+
+
         /// <summary>
         /// Get the list of clauses for a component.
         /// </summary>
@@ -123,13 +145,10 @@ public List<AbstractClause> GetComponents(string component, string engineCode =
         /// <returns></returns>
         public C GetOneComponent<C>(string component, string engineCode = null) where C : AbstractClause
         {
-            if (engineCode == null)
-            {
-                engineCode = EngineScope;
-            }
+            engineCode = engineCode ?? EngineScope;
 
-            return GetComponents<C>(component, engineCode)
-            .FirstOrDefault();
+            var all = GetComponents<C>(component, engineCode);
+            return all.FirstOrDefault(c => c.Engine == engineCode) ?? all.FirstOrDefault(c => c.Engine == null);
         }
 
         /// <summary>
@@ -149,7 +168,7 @@ public AbstractClause GetOneComponent(string component, string engineCode = null
         }
 
         /// <summary>
-        /// Return wether the query has clauses for a component.
+        /// Return whether the query has clauses for a component.
         /// </summary>
         /// <param name="component"></param>
         /// <param name="engineCode"></param>
@@ -247,9 +266,9 @@ protected bool GetNot()
         /// <returns></returns>
         public Q From(string table)
         {
-            return ClearComponent("from").AddComponent("from", new FromClause
+            return AddOrReplaceComponent("from", new FromClause
             {
-                Table = table
+                Table = table,
             });
         }
 
@@ -263,15 +282,15 @@ public Q From(Query query, string alias = null)
                 query.As(alias);
             };
 
-            return ClearComponent("from").AddComponent("from", new QueryFromClause
+            return AddOrReplaceComponent("from", new QueryFromClause
             {
                 Query = query
             });
         }
 
         public Q FromRaw(string sql, params object[] bindings)
         {
-            return ClearComponent("from").AddComponent("from", new RawFromClause
+            return AddOrReplaceComponent("from", new RawFromClause
             {
                 Expression = sql,
                 Bindings = bindings,