Merge 3b4007a into e11a354

docs/components.rst

@@ -27,6 +27,7 @@ Currently, there are three different kinds of commands:
         * *SendMail*, 
         * *Download*,
         * *CopyFromFile*,
+        * *CopyToFile*,
         * *Shutdown*.
 
 Task
@@ -156,6 +157,15 @@ Kind
                 "filename": "download/orte_ansi.txt" 
             }'::jsonb
 
+``BUILTIN: CopyToFile``
+    ``object``
+        .. code-block:: SQL
+                
+            '{
+                "sql": "COPY location TO STDOUT", 
+                "filename": "download/location.txt" 
+            }'::jsonb
+
 ``BUILTIN: Shutdown``
     *value ignored*
 

internal/pgengine/copy.go

@@ -5,6 +5,22 @@ import (
 	"os"
 )
 
+// CopyToFile copies data from database into local file using COPY format specified by sql
+func (pge *PgEngine) CopyToFile(ctx context.Context, filename string, sql string) (int64, error) {
+	dbconn, err := pge.ConfigDb.Acquire(ctx)
+	if err != nil {
+		return -1, err
+	}
+	defer dbconn.Release()
+	f, err := os.Create(filename)
+	defer func() { _ = f.Close() }()
+	if err != nil {
+		return -1, err
+	}
+	res, err := dbconn.Conn().PgConn().CopyTo(ctx, f, sql)
+	return res.RowsAffected(), err
+}
+
 // CopyFromFile copies data from local file into database using COPY format specified by sql
 func (pge *PgEngine) CopyFromFile(ctx context.Context, filename string, sql string) (int64, error) {
 	dbconn, err := pge.ConfigDb.Acquire(ctx)

internal/pgengine/copy_test.go

@@ -22,3 +22,15 @@ func TestCopyFromFile(t *testing.T) {
 	assert.True(t, cnt == 2, "Should copy exactly 2 rows")
 	assert.NoError(t, os.RemoveAll("test.csv"), "Test output should be removed")
 }
+
+func TestCopyToFile(t *testing.T) {
+	teardownTestCase := SetupTestCase(t)
+	defer teardownTestCase(t)
+	ctx := context.Background()
+	_, err := pge.CopyToFile(ctx, "", "COPY location TO STDOUT")
+	assert.Error(t, err, "Should fail for empty file name")
+	cnt, err := pge.CopyToFile(ctx, "test.csv", "COPY (SELECT generate_series(1,5)) TO STDOUT (FORMAT csv)")
+	assert.NoError(t, err, "Should copy to file")
+	assert.True(t, cnt == 5, "Should copy exactly 5 rows")
+	assert.NoError(t, os.RemoveAll("test.csv"), "Test output should be removed")
+}

internal/scheduler/tasks.go

@@ -20,6 +20,7 @@ var Tasks = map[string](func(context.Context, *Scheduler, string) (string, error
 	"SendMail":     taskSendMail,
 	"Download":     taskDownload,
 	"CopyFromFile": taskCopyFromFile,
+	"CopyToFile":   taskCopyToFile,
 	"Shutdown":     taskShutdown}
 
 func (sch *Scheduler) executeTask(ctx context.Context, name string, paramValues []string) (stdout string, err error) {
@@ -85,6 +86,22 @@ func taskCopyFromFile(ctx context.Context, sch *Scheduler, val string) (stdout s
 	return stdout, err
 }
 
+func taskCopyToFile(ctx context.Context, sch *Scheduler, val string) (stdout string, err error) {
+	type copyTo struct {
+		SQL      string `json:"sql"`
+		Filename string `json:"filename"`
+	}
+	var ct copyTo
+	if err := json.Unmarshal([]byte(val), &ct); err != nil {
+		return "", err
+	}
+	count, err := sch.pgengine.CopyToFile(ctx, ct.Filename, ct.SQL)
+	if err == nil {
+		stdout = fmt.Sprintf("%d rows copied to %s", count, ct.Filename)
+	}
+	return stdout, err
+}
+
 func taskDownload(ctx context.Context, sch *Scheduler, paramValues string) (stdout string, err error) {
 	type downloadOpts struct {
 		WorkersNum int      `json:"workersnum"`

internal/scheduler/tasks_test.go

@@ -35,6 +35,9 @@ func TestExecuteTask(t *testing.T) {
 	assert.Error(t, et("CopyFromFile", []string{"foo"}), "Invalid json")
 	assert.Error(t, et("CopyFromFile", []string{`{"sql": "COPY", "filename": "foo"}`}), "Acquire() should fail")
 
+	assert.Error(t, et("CopyToFile", []string{"foo"}), "Invalid json")
+	assert.Error(t, et("CopyToFile", []string{`{"sql": "COPY", "filename": "foo"}`}), "Acquire() should fail")
+
 	assert.Error(t, et("SendMail", []string{"foo"}), "Invalid json")
 	assert.Error(t, et("SendMail", []string{`{"ServerHost":"smtp.example.com","ServerPort":587,"Username":"user"}`}))
 
@@ -44,4 +47,5 @@ func TestExecuteTask(t *testing.T) {
 	assert.Error(t, et("Download", []string{`{"workersnum": 0, "fileurls": ["http://foo.bar"], "destpath": "" }`}),
 		"Downlod incorrect url should fail")
 
+	assert.NoError(t, et("Shutdown", []string{}))
 }