add copy command support to Npgsql (#72)

DaredevilOSS · Jul 20, 2024 · 3683023 · 3683023
1 parent 077e21a
commit 3683023
Show file tree

Hide file tree

Showing 21 changed files with 250 additions and 100 deletions.
diff --git a/CodeGenerator/CodeGenerator.cs b/CodeGenerator/CodeGenerator.cs
@@ -238,7 +238,10 @@ private MemberDeclarationSyntax AddMethodDeclaration(Query query)
                 query.Params, query.Columns),
             ":many" => DbDriver.ManyDeclare(query.Name, queryTextConstant, argInterface, returnInterface,
                 query.Params, query.Columns),
-            ":execlastid" => DbDriver.ExecLastIdDeclare(query.Name, queryTextConstant, argInterface, query.Params),
+            ":execlastid" => ((MySqlConnectorDriver)DbDriver)
+                .ExecLastIdDeclare(query.Name, queryTextConstant, argInterface, query.Params),
+            ":copyfrom" => ((NpgsqlDriver)DbDriver)
+                .CopyFromDeclare(query.Name, queryTextConstant, argInterface, query.Params),
             _ => throw new InvalidDataException()
         };
 

diff --git a/Drivers/DbDriver.cs b/Drivers/DbDriver.cs
@@ -59,7 +59,7 @@ public string GetColumnReader(Column column, int ordinal)
 
     public abstract string TransformQueryText(Query query);
 
-    public abstract (string, string) EstablishConnection();
+    public abstract (string, string) EstablishConnection(bool isCopyCommand = false); // TODO fix codesmell - should act upon the query object
 
     public abstract string CreateSqlCommand(string sqlTextConstant);
 
@@ -72,9 +72,6 @@ public abstract MemberDeclarationSyntax ManyDeclare(string funcName, string sqlT
     public abstract MemberDeclarationSyntax ExecDeclare(string funcName, string text, string argInterface,
         IList<Parameter> parameters);
 
-    public abstract MemberDeclarationSyntax ExecLastIdDeclare(string funcName, string queryTextConstant,
-        string argInterface, IList<Parameter> parameters);
-
     public static bool IsCsharpPrimitive(string csharpType)
     {
         var csharpPrimitives = new HashSet<string> { "long", "double", "int", "float", "bool" };

diff --git a/Drivers/Generators/CopyFromDeclareGen.cs b/Drivers/Generators/CopyFromDeclareGen.cs
@@ -0,0 +1,83 @@
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Plugin;
+using System.Collections.Generic;
+using System.Linq;
+using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
+
+namespace SqlcGenCsharp.Drivers.Generators;
+
+public class CopyFromDeclareGen(DbDriver dbDriver)
+{
+    public MemberDeclarationSyntax Generate(string funcName, string queryTextConstant, string argInterface,
+        IList<Parameter> parameters)
+    {
+        return ParseMemberDeclaration($$"""
+                                        public async Task {{funcName}}(List<{{argInterface}}> args)
+                                        {
+                                            {{GetMethodBody(queryTextConstant, parameters)}}
+                                        }
+                                        """)!;
+    }
+
+    private string GetMethodBody(string queryTextConstant, IEnumerable<Parameter> parameters)
+    {
+        var (establishConnection, connectionOpen) = dbDriver.EstablishConnection(true);
+        var beginBinaryImport = $"{Variable.Connection.Name()}.BeginBinaryImportAsync({queryTextConstant}";
+
+        return dbDriver.DotnetFramework.LatestDotnetSupported()
+            ? GetAsModernDotnet()
+            : GetAsLegacyDotnet();
+
+        string GetAsModernDotnet()
+        {
+            var addRowsToCopyCommand = AddRowsToCopyCommand();
+            return $$"""
+                     {
+                         await using {{establishConnection}};
+                         {{connectionOpen.AppendSemicolonUnlessEmpty()}}
+                         await {{Variable.Connection.Name()}}.OpenAsync();
+                         await using var {{Variable.Writer.Name()}} = await {{beginBinaryImport}});
+                         {{addRowsToCopyCommand}}
+                         await {{Variable.Writer.Name()}}.CompleteAsync();
+                         await {{Variable.Connection.Name()}}.CloseAsync();
+                     }
+                     """;
+        }
+
+        string GetAsLegacyDotnet()
+        {
+            var addRowsToCopyCommand = AddRowsToCopyCommand();
+            return $$"""
+                     {
+                         using ({{establishConnection}})
+                         {
+                             {{connectionOpen.AppendSemicolonUnlessEmpty()}}
+                             await {{Variable.Connection.Name()}}.OpenAsync();
+                             using (var {{Variable.Writer.Name()}} = await {{beginBinaryImport}}))
+                             {
+                                {{addRowsToCopyCommand}}
+                                await {{Variable.Writer.Name()}}.CompleteAsync();
+                             }
+                             await {{Variable.Connection.Name()}}.CloseAsync();
+                         }
+                     }
+                     """;
+        }
+
+        string AddRowsToCopyCommand()
+        {
+            var constructRow = new List<string>()
+                    .Append($"await {Variable.Writer.Name()}.StartRowAsync();")
+                    .Concat(parameters
+                        .Select(p =>
+                            $"await {Variable.Writer.Name()}.WriteAsync({Variable.Row.Name()}.{p.Column.Name.FirstCharToUpper()});"))
+                    .JoinByNewLine();
+            return $$"""
+                   foreach (var {{Variable.Row.Name()}} in args) 
+                   {
+                        {{constructRow}}
+                   }
+                   """;
+        }
+    }
+}
diff --git a/Drivers/Generators/ExecDeclareGen.cs b/Drivers/Generators/ExecDeclareGen.cs
@@ -10,8 +10,9 @@ public class ExecDeclareGen(DbDriver dbDriver)
     public MemberDeclarationSyntax Generate(string funcName, string queryTextConstant, string argInterface,
         IList<Parameter> parameters)
     {
+        var parametersStr = CommonGen.GetParameterListAsString(argInterface, parameters);
         return ParseMemberDeclaration($$"""
-                                        public async Task {{funcName}}({{CommonGen.GetParameterListAsString(argInterface, parameters)}})
+                                        public async Task {{funcName}}({{parametersStr}})
                                         {
                                             {{GetMethodBody(queryTextConstant, parameters)}}
                                         }
@@ -34,7 +35,7 @@ string GetWithUsingAsStatement()
             return $$"""
                      {
                          await using {{establishConnection}};
-                         {{connectionOpen}};
+                         {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                          await using {{createSqlCommand}};
                          {{commandParameters.JoinByNewLine()}}
                          {{executeScalar}}
@@ -48,7 +49,7 @@ string GetWithUsingAsBlock()
                      {
                          using ({{establishConnection}})
                          {
-                             {{connectionOpen}};
+                             {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                              using ({{createSqlCommand}})
                              {
                                  {{commandParameters.JoinByNewLine()}}

diff --git a/Drivers/Generators/ExecLastIdDeclareGen.cs b/Drivers/Generators/ExecLastIdDeclareGen.cs
@@ -10,8 +10,9 @@ public class ExecLastIdDeclareGen(DbDriver dbDriver)
     public MemberDeclarationSyntax Generate(string funcName, string queryTextConstant, string argInterface,
         IList<Parameter> parameters)
     {
+        var parametersStr = CommonGen.GetParameterListAsString(argInterface, parameters);
         return ParseMemberDeclaration($$"""
-                                        public async Task<long> {{funcName}}({{CommonGen.GetParameterListAsString(argInterface, parameters)}})
+                                        public async Task<long> {{funcName}}({{parametersStr}})
                                         {
                                             {{GetMethodBody(queryTextConstant, parameters)}}
                                         }
@@ -34,7 +35,7 @@ string GetWithUsingAsStatement()
             return $$"""
                      {
                          await using {{establishConnection}};
-                         {{connectionOpen}};
+                         {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                          await using {{createSqlCommand}};
                          {{commandParameters.JoinByNewLine()}}
                          {{executeScalarAndReturnCreated.JoinByNewLine()}}
@@ -48,7 +49,7 @@ string GetWithUsingAsBlock()
                      {
                          using ({{establishConnection}})
                          {
-                             {{connectionOpen}};
+                             {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                              using ({{createSqlCommand}})
                              {
                                 {{commandParameters.JoinByNewLine()}}

diff --git a/Drivers/Generators/ManyDeclareGen.cs b/Drivers/Generators/ManyDeclareGen.cs
@@ -13,10 +13,10 @@ public class ManyDeclareGen(DbDriver dbDriver)
     public MemberDeclarationSyntax Generate(string funcName, string queryTextConstant, string argInterface,
         string returnInterface, IList<Parameter> parameters, IEnumerable<Column> columns)
     {
-        var parameterList = CommonGen.GetParameterListAsString(argInterface, parameters);
+        var parametersStr = CommonGen.GetParameterListAsString(argInterface, parameters);
         var returnType = $"Task<List<{returnInterface}>>";
         return ParseMemberDeclaration($$"""
-                                        public async {{returnType}} {{funcName}}({{parameterList}})
+                                        public async {{returnType}} {{funcName}}({{parametersStr}})
                                         {
                                             {{GetMethodBody(queryTextConstant, returnInterface, columns, parameters)}}
                                         }
@@ -48,7 +48,7 @@ string GetWithUsingAsStatement()
             return $$"""
                      {
                          await using {{establishConnection}};
-                         {{connectionOpen}};
+                         {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                          await using {{createSqlCommand}};
                          {{commandParameters.JoinByNewLine()}}
                          {{initDataReader}};
@@ -65,7 +65,7 @@ string GetWithUsingAsBlock()
                      {
                          using ({{establishConnection}})
                          {
-                             {{connectionOpen}};
+                             {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                              using ({{createSqlCommand}})
                              {
                                  {{commandParameters.JoinByNewLine()}}

diff --git a/Drivers/Generators/OneDeclareGen.cs b/Drivers/Generators/OneDeclareGen.cs
@@ -42,7 +42,7 @@ string GetWithUsingAsStatement()
             return $$"""
                      {
                          await using {{establishConnection}};
-                         {{connectionOpen}};
+                         {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                          await using {{createSqlCommand}};
                          {{commandParameters.JoinByNewLine()}}
                          {{initDataReader}};
@@ -61,7 +61,7 @@ string GetWithUsingAsBlock()
                      {
                          using ({{establishConnection}})
                          {
-                             {{connectionOpen}};
+                             {{connectionOpen.AppendSemicolonUnlessEmpty()}}
                              using ({{createSqlCommand}})
                              {
                                 {{commandParameters.JoinByNewLine()}}

diff --git a/Drivers/MySqlConnectorDriver.cs b/Drivers/MySqlConnectorDriver.cs
@@ -58,7 +58,7 @@ public override UsingDirectiveSyntax[] GetUsingDirectives()
             .ToArray();
     }
 
-    public override (string, string) EstablishConnection()
+    public override (string, string) EstablishConnection(bool isCopyCommand = false)
     {
         return (
             $"var {Variable.Connection.Name()} = new MySqlConnection({Variable.ConnectionString.Name()})",
@@ -89,7 +89,7 @@ public override MemberDeclarationSyntax ExecDeclare(string funcName, string quer
         return new ExecDeclareGen(this).Generate(funcName, queryTextConstant, argInterface, parameters);
     }
 
-    public override MemberDeclarationSyntax ExecLastIdDeclare(string funcName, string queryTextConstant,
+    public MemberDeclarationSyntax ExecLastIdDeclare(string funcName, string queryTextConstant,
         string argInterface, IList<Parameter> parameters)
     {
         return new ExecLastIdDeclareGen(this).Generate(funcName, queryTextConstant, argInterface, parameters);

diff --git a/Drivers/NpgsqlDriver.cs b/Drivers/NpgsqlDriver.cs
@@ -47,11 +47,14 @@ public override UsingDirectiveSyntax[] GetUsingDirectives()
             .ToArray();
     }
 
-    public override (string, string) EstablishConnection()
+    public override (string, string) EstablishConnection(bool isCopyCommand = false)
     {
-        return (
-            $"var {Variable.Connection.Name()} = NpgsqlDataSource.Create({Variable.ConnectionString.Name()})",
-            string.Empty);
+        if (isCopyCommand)
+            return (
+                $"var ds = NpgsqlDataSource.Create({Variable.ConnectionString.Name()})",
+                $"var {Variable.Connection.Name()} = ds.CreateConnection()"
+            );
+        return ($"var {Variable.Connection.Name()} = NpgsqlDataSource.Create({Variable.ConnectionString.Name()})", "");
     }
 
     public override string CreateSqlCommand(string sqlTextConstant)
@@ -61,15 +64,23 @@ public override string CreateSqlCommand(string sqlTextConstant)
 
     public override string TransformQueryText(Query query)
     {
+        if (query.Cmd == ":copyfrom")
+            return GetCopyCommand();
+
         var queryText = query.Text;
         for (var i = 0; i < query.Params.Count; i++)
         {
             var currentParameter = query.Params[i];
             queryText = Regex.Replace(queryText, $@"\$\s*{i + 1}",
                 $"@{currentParameter.Column.Name.FirstCharToLower()}");
         }
-
         return queryText;
+
+        string GetCopyCommand()
+        {
+            var copyParams = query.Params.Select(p => p.Column.Name).JoinByComma();
+            return $"COPY {query.InsertIntoTable.Name} ({copyParams}) FROM STDIN (FORMAT BINARY)";
+        }
     }
 
     public override MemberDeclarationSyntax OneDeclare(string funcName, string queryTextConstant, string argInterface,
@@ -85,16 +96,16 @@ public override MemberDeclarationSyntax ExecDeclare(string funcName, string quer
         return new ExecDeclareGen(this).Generate(funcName, queryTextConstant, argInterface, parameters);
     }
 
-    public override MemberDeclarationSyntax ExecLastIdDeclare(string funcName, string queryTextConstant,
-        string argInterface, IList<Parameter> parameters)
-    {
-        return new ExecLastIdDeclareGen(this).Generate(funcName, queryTextConstant, argInterface, parameters);
-    }
-
     public override MemberDeclarationSyntax ManyDeclare(string funcName, string queryTextConstant, string argInterface,
         string returnInterface, IList<Parameter> parameters, IEnumerable<Column> columns)
     {
         return new ManyDeclareGen(this).Generate(funcName, queryTextConstant, argInterface, returnInterface, parameters,
             columns);
     }
+
+    public MemberDeclarationSyntax CopyFromDeclare(string funcName, string queryTextConstant, string argInterface,
+        IList<Parameter> parameters)
+    {
+        return new CopyFromDeclareGen(this).Generate(funcName, queryTextConstant, argInterface, parameters);
+    }
 }
diff --git a/Drivers/Variable.cs b/Drivers/Variable.cs
@@ -5,6 +5,8 @@ public enum Variable
     ConnectionString,
     Connection,
     Reader,
+    Row,
+    Writer,
     Command,
     Result
 }

diff --git a/EndToEndTests/Consts.cs → EndToEndTests/DataGenerator.cs b/EndToEndTests/Consts.cs → EndToEndTests/DataGenerator.cs
@@ -1,6 +1,9 @@
+using NpgsqlExample;
+using System;
+
 namespace SqlcGenCsharpTests;
 
-public static class Consts
+public static class DataGenerator
 {
     public const string BojackAuthor = "Bojack Horseman";
     public const string BojackTheme = "Back in the 90s he was in a very famous TV show";

diff --git a/EndToEndTests/ISqlDriverTester.cs b/EndToEndTests/ISqlDriverTester.cs
@@ -2,18 +2,18 @@
 
 namespace SqlcGenCsharpTests;
 
-public interface ISqlDriverTester
+public abstract class SqlDriverTester
 {
-    async Task TestFlow()
+    public async Task TestBasicFlow()
     {
         var firstInsertedId = await CreateFirstAuthorAndTest();
         await CreateSecondAuthorAndTest();
         await DeleteFirstAuthorAndTest(firstInsertedId);
     }
 
-    protected Task<long> CreateFirstAuthorAndTest();
+    protected abstract Task<long> CreateFirstAuthorAndTest();
 
-    protected Task CreateSecondAuthorAndTest();
+    protected abstract Task CreateSecondAuthorAndTest();
 
-    protected Task DeleteFirstAuthorAndTest(long idToDelete);
+    protected abstract Task DeleteFirstAuthorAndTest(long idToDelete);
 }