如何从STDIN导入Postgresql的行?

时间:2016-02-13 18:55:22

标签: postgresql go

在Python中,我有以下内容,可以在不使用文件的情况下将行批量加载到Postgresql:

import csv
import subprocess

mylist, keys = [{'name': 'fred'}, {'name': 'mary'}], ['name']
p = subprocess.Popen(['psql', 'mydb', '-U', 'openupitsme', '-h', 'my.ip.address', '--no-password', '-c',
    '\COPY tester(%s) FROM STDIN (FORMAT CSV)' % ', '.join(keys),
    '--set=ON_ERROR_STOP=false'
    ], stdin=subprocess.PIPE
)
for d in mylist:
    dict_writer = csv.DictWriter(p.stdin, keys, quoting=csv.QUOTE_MINIMAL)
    dict_writer.writerow(d)
p.stdin.close()

我想在Go中完成同样的事情。我目前正在将行写入文件,然后导入它们然后删除该文件。我想像在Python中一样从STDIN导入行。我有:

package main

import (
    "database/sql"
    "log"
    "os"
    "os/exec"

    _ "github.com/lib/pq"
)

var (
    err error
    db  *sql.DB
)

func main() {
    var err error
    fh := "/path/to/my/file.txt"
    f, err := os.Create(fh)
    if err != nil {
        panic(err)
    }
    defer f.Close()
    defer os.Remove(fh)
    rows := []string{"fred", "mary"}
    for _, n := range rows {
        _, err = f.WriteString(n + "\n")
        if err != nil {
            panic(err)
        }
    }
    // dump to postgresql
    c := exec.Command("psql", "mydb", "-U", "openupitsme", "-h", "my.ip.address", "--no-password",
        "-c", `\COPY tester(customer) FROM `+fh)
    if out, err := c.CombinedOutput(); err != nil {
        log.Println(string(out), err)
    }
}

编辑: 更进一步,但这不是插入记录:

    keys := []string{"link", "domain"}
    records := [][]string{
        {"first_name", "last_name"},
        {"Rob", "Pike"},
        {"Ken", "Thompson"},
        {"Robert", "Griesemer"},
    }

    cmd := exec.Command("psql")
    stdin, err := cmd.StdinPipe()
    if err != nil {
        log.Println(err)
    }
    stdout, err := cmd.StdoutPipe()
    if err != nil {
        log.Println(err)
    }
    if err := cmd.Start(); err != nil {
        log.Println(err)
    }
    go func() {
        _, err = io.WriteString(stdin, "search -U meyo -h 1.2.3.4 -p 1111 --no-password -c ")
        if err != nil {
            log.Println(err)
        }
        _, err := io.WriteString(stdin, fmt.Sprintf("COPY links(%s) FROM STDIN (FORMAT CSV)", strings.Join(keys, ",")))
        if err != nil {
            log.Println(err)
        }
        w := csv.NewWriter(stdin)
        if err := w.WriteAll(records); err != nil {
            log.Fatalln("error writing record to csv:", err)
        }
        w.Flush()
        if err := w.Error(); err != nil {
            log.Fatal(err)
        }
        if err != nil {
            log.Println(err)
        }
        stdin.Close()
    }()

    done := make(chan bool)
    go func() {
        _, err := io.Copy(os.Stdout, stdout)
        if err != nil {
            log.Fatal(err)
        }
        stdout.Close()
        done <- true
    }()
    <-done

    if err := cmd.Wait(); err != nil {
        log.Println(err, cmd.Args, stdout)
    }

没有插入记录,我收到一个无用的错误:

exit status 2

2 个答案:

答案 0 :(得分:1)

以下代码应指向您想要的方向:

package main

import (
    "fmt"
    "log"
    "os"
    "os/exec"
    "strings"
)

func main() {
    keys := []string{"customer"}
    sqlCmd := fmt.Sprintf("COPY tester(%s) FROM STDIN (FORMAT CSV)", strings.Join(keys, ","))
    cmd := exec.Command("psql", "<dbname>", "-U", "<username>", "-h", "<host_ip>", "--no-password", "-c", sqlCmd)
    cmd.Stdin = os.Stdin
    output, _ := cmd.CombinedOutput()
    log.Println(string(output))
}

如果密钥需要是动态的,您可以从os.Args收集它们。

请注意,如果您打算使用psql命令,则无需导入database / sql或lib / pq。如果您对使用lib / pq感兴趣,请查看lib / pq文档中的Bulk Imports

答案 1 :(得分:1)

github.com/lib/pq软件包文档实际上有an example如何做你想要的。以下是整个计划的改编文本:

package main

import (
    "database/sql"
    "log"

    "github.com/lib/pq"
)

func main() {
    records := [][]string{
        {"Rob", "Pike"},
        {"Ken", "Thompson"},
        {"Robert", "Griesemer"},
    }

    db, err := sql.Open("postgres", "dbname=postgres user=postgres password=postgres")
    if err != nil {
        log.Fatalf("open: %v", err)
    }
    if err = db.Ping(); err != nil {
        log.Fatalf("open ping: %v", err)
    }
    defer db.Close()

    txn, err := db.Begin()
    if err != nil {
        log.Fatalf("begin: %v", err)
    }

    stmt, err := txn.Prepare(pq.CopyIn("test", "first_name", "last_name"))
    if err != nil {
        log.Fatalf("prepare: %v", err)
    }

    for _, r := range records {
        _, err = stmt.Exec(r[0], r[1])
        if err != nil {
            log.Fatalf("exec: %v", err)
        }
    }

    _, err = stmt.Exec()
    if err != nil {
        log.Fatalf("exec: %v", err)
    }

    err = stmt.Close()
    if err != nil {
        log.Fatalf("stmt close: %v", err)
    }

    err = txn.Commit()
    if err != nil {
        log.Fatalf("commit: %v", err)
    }
}

在我的机器上,这会在大约2秒内导入1 000 000条记录。