golang解析csv并导入elastic search

今天我们一起来写一个从csv文件将数据导入elastic search的小程序

准备工作

在gopath的src文件夹下创建csv2es文件夹，并创建main.go文件。下载一些csv文件备用

解析命令行参数

首先我们需要使用flag package解析命令行参数,代码如下

func main() {
  // 解析命令行输入
  host := flag.String("host", "http://localhost:9200", "host, e.g. http://localhost:9200")
  file := flag.String("file", "", "file path")
  esIndex := flag.String("index", "", "elastic search index")
  esType := flag.String("type", "", "elastic search type")
  flag.Parse()
  if *file == "" {
    fmt.Println("please set which csv file you want to import clearly")
    return
  }
  if *esIndex == "" {
    fmt.Println("please set elastic search index")
    return
  }
  if *esType == "" {
    fmt.Println("please set elastic search type")
    return
  }
...
}

主要解析的参数有es的地址，待导入的文件的路径，导入到的es的index和type。运行go build，生成可执行文件csv2go,执行

csv2go -h
-file string
        file path (default "")
  -host string
        host, e.g. http://localhost:9200 (default "http://localhost:9200")
  -index string
        elastic search index (default "")
  -type string
        elastic search type (default "")

连接es

elastic这个开源项目可以帮助我们连接elastic

  // 连接es
  ctx := context.Background()
  client, err := elastic.NewClient(
    elastic.SetURL(*host),
    elastic.SetSniff(false))
  if err != nil {
    panic(err)
  }

  // 检查index是否存在，如果不存在则创建index
  exists, err := client.IndexExists(*esIndex).Do(ctx)
  if err != nil {
    panic(err)
  }
  if !exists {
    createIndex, err := client.CreateIndex(*esIndex).Do(ctx)
    if err != nil {
      panic(err)
    }
  }

解析csv并导入(index)到elastic search

这里需要注意几点。第一，Mac上会存在\r结尾的文件的问题，所以我们使用macreader这个包对io.Reader包了一层，有兴趣的同学可以看我之前的文章《mac上的文件有毒》，第二，我们默认csv文件的第一行为column name，后面各行都是合法的记录。

  // 解析csv
  f, _ := os.Open(*file)
  r := csv.NewReader(macreader.New(bufio.NewReader(f)))
  keys, err := r.Read()
  for {                                            //1
    record, err := r.Read()
    if err == io.EOF {
      break
    }
    m := make(map[string]string)
    for i, key := range keys {
      m[key] = record[i]
    }
    jsonStr, err := json.Marshal(m)
    if err != nil {
      panic(err)
    }
    put1, err := client.Index().
            Index(*esIndex).
            Type(*esType).
            BodyString(string(jsonStr)).
            Do(ctx)
    if err != nil {
      // Handle error
      panic(err)
    }
    fmt.Printf("Indexed tweet %s to index %s, type %s\n", put1.Id, put1.Index, put1.Type)
  }                                           //2

ok,一个基本的将csv中的数据导入elastic search的程序完成了，我们来测试一下性能吧。在上面的（1）行代码前面加上

start := time.Now().Unix()

在上面的（2）行代码后面加上

end := time.Now().Unix()
fmt.Println(end, start)

测试了一下跑了一个61567条记录的文件，一共跑了36分钟。够我睡一个午觉了...

提升效率

elastic search有一个bulk api,可以将一些操作合并起来，同时传递给elastic search处理并返回

  // 新建一个Bulk
  bulkRequest := client.Bulk()
  for {
    ...
    // 为每一条记录生成一个IndexRequest并加入Bulk
    req := elastic.NewBulkIndexRequest().Index(*esIndex).Type(*esType).Doc(string(jsonStr))
    bulkRequest.Add(req)
  }
  // 一次性完成请求
  bulkResponse, err := bulkRequest.Do(ctx)
  if err != nil {
  }
  indexed := bulkResponse.Indexed()
  fmt.Println("向es导入了",len(indexed),"条数据")

优化过后，插入相同的6万多条记录只需要几秒钟。cool ^_

本文的代码已经开源在github,欢迎使用或者提出意见。