Go言語でCSVファイルのデータをPostgreSQLに簡単に保存する方法

はじめに

Go言語を使ってCSVファイルからデータを読み込み、PostgreSQLデータベースに保存する方法を紹介します。この記事では、コマンドライン引数を使用してCSVファイルのパスを指定し、データベースにデータを保存する手順を詳しく説明します。

なお、こちら以下の記事の続きになります。

csvファイルの準備

dataディレクトリの配下にarticlesディクレトリを作成し、その中にmainichi.csvを入れました。

data/articles/mainichi.csv

2024,12,21,大正デモクラシーが高揚する中....
2024,12,22,漫画家、水島新司さん（故人）....
2024,12,23,言い得て妙とは、このことか。....

このように、年、月、日、本文の順番になっています。

バッチ処理（本文挿入のスクリプト）の作成

ファイルは以下のようにしました。

batch.go

package main

import (
    "context"
    "encoding/csv"
    "flag"
    "log"
    "os"
	"fmt"
    "strconv"
	"github.com/joho/godotenv"

    "github.com/jijimama/newspaper-app/ent"
	"github.com/jijimama/newspaper-app/ent/newspaper"
    _ "github.com/lib/pq"
)

func main() {
    // コマンドライン引数を解析
    newspaperName := flag.String("newspaper", "", "Name of the newspaper")
    csvFilePath := flag.String("csv", "data/articles/mainichi.csv", "Path to the CSV file")
    flag.Parse()

    // 新聞名が指定されていない場合はエラーを出力して終了
    if *newspaperName == "" {
        log.Fatal("newspaper name is required")
    }

    // .envファイルを読み込む
	err := godotenv.Load()
	if err != nil {
		log.Fatalf("Error loading .env file: %v", err)
	}

	// 環境変数から取得
	dbHost := os.Getenv("DB_HOST")
	dbPort := os.Getenv("DB_PORT")
	dbUser := os.Getenv("DB_USER")
	dbName := os.Getenv("DB_NAME")
	dbPassword := os.Getenv("DB_PASSWORD")

	// 接続文字列を作成
	dsn := "host=" + dbHost + " port=" + dbPort + " user=" + dbUser + " dbname=" + dbName + " password=" + dbPassword + " sslmode=disable"

    //PostgreSQLに接続
	client, err := ent.Open("postgres", dsn)
	if err != nil {
		log.Fatalf("failed opening connection to postgres: %v", err)
	}
	defer client.Close()
	// Run the auto migration tool.
	if err := client.Schema.Create(context.Background()); err != nil {
		log.Fatalf("failed creating schema resources: %v", err)
	}

    // 指定された新聞名で新聞レコードをクエリ
    newspaper, err := client.Newspaper.
        Query().
        Where(newspaper.NameEQ(*newspaperName)).
        Only(context.Background())
    if err != nil {
        log.Fatalf("failed searching newspaper: %v", err)
    }

    // CSVファイルからデータを読み込み、データベースに挿入
    err = LoadCSVData(client, *csvFilePath, newspaper)
    if err != nil {
        log.Fatalf("failed loading CSV data: %v", err)
    }
}

// LoadCSVData は、指定されたCSVファイルからデータを読み込み、データベースに挿入します。
func LoadCSVData(client *ent.Client, filePath string, newspaper *ent.Newspaper) error {
    // CSVファイルを開く
    file, err := os.Open(filePath)
    if err != nil {
        return err
    }
    defer file.Close()

    // CSVリーダーを作成
    reader := csv.NewReader(file)
    records, err := reader.ReadAll()
    if err != nil {
        return err
    }

    // 各レコードを処理
    for _, record := range records {
        // 年、月、日を整数に変換
        year, err := strconv.Atoi(record[0])
        if err != nil {
            return err
        }
        month, err := strconv.Atoi(record[1])
        if err != nil {
            return err
        }
        day, err := strconv.Atoi(record[2])
        if err != nil {
            return err
        }
        content := record[3]

        // 新しい記事レコードを作成し、データベースに保存
        _, err = client.Article.
            Create().
            SetYear(year).
            SetMonth(month).
            SetDay(day).
            SetContent(content).
            SetNewspaper(newspaper).
            Save(context.Background())
        if err != nil {
            return err
        }
    }

    return nil
}

バッチ処理（本文挿入のスクリプト）の実行

docker compose exec backend bash
# go run batch.go -newspaper="毎日新聞" -csv="data/articles/mainichi.csv"

newspaperとcsvファイルの指定を行います。

データの確認

docker compose exec db bash
psql -U postgres -d db
\x 
db=# SELECT * FROM articles;
-[ RECORD 1 ]------+
id                 | 1
content            | 大正デモクラシーが高揚する中、.....
year               | 2024
month              | 12
day                | 21
created_at         | 2024-12-24 05:10:36.373939+00
updated_at         | 2024-12-24 05:10:36.373939+00
newspaper_articles | 1

無事にデータ挿入ができました！