본문 바로가기

Lang

[go] postgresql 연동 테스트용 go 코딩 샘플

외부 서버에 postgresql 설치 후 원격 접속 확인 하기 위해 프로그램 급조.

원격 접속 확인하기 위해 DB 클라이언트 툴 설치한다거나 테스트용 프로그램 실행 환경을 추가 설치(가령 jre 같은거)한다는게 어려운 여건이어서 근 이 년여 넘게 안써서 기초 문법 조차 다 잊어먹은 go 언어를 간만에 소환.

소스 빌드해서 실행 파일 하나와 db 정보 설정 위한 yaml 파일, 이렇게만 배포하면 되도록 구생했고 접속 확인 및 권한 부여 제대로 되었나 확인하는 차원에서 임시 테이블 생성, insert, select 후 결과 출력, 임시 테이블 drop 순으로 동작하도록 구현.

온라인에서 필요한 기능들의 샘플 코드 찾아 조합한거고 구문도 설명 필요 없을 정도로 간단한지라 코딩했다 하기 민망하지만 애매한 테스트 환경에서 이런 식의 아이디어 응용해볼 수 있지 않을까 싶어서 내부 공유했던거 살짝 편집해서 올려본다.

yaml, postgresql 사용하는 매우 기초적인 구문들만 사용했으니 조금씩만 수정하면 클라이언트 툴 설치 등이 여의치 않은 환경에서 간단한 db 작업에 응용할 수 있지 않을까나?

// dbtest.yaml

host: ....
port: ...
database: ...
user: ...
password: ...

// go 소스

// go 문법 몰라도 그냥 이해할 수 있는 정도라 구문 설명은 스킵.

package main

import (
    "database/sql"
    "fmt"
    _ "github.com/lib/pq"
    "gopkg.in/yaml.v3"
    "io/ioutil"
    "log"
)

type Pdb struct {
    dbObj *sql.DB
    table string
}

func (db *Pdb) DbObj() *sql.DB {
    return db.dbObj
}

func (db *Pdb) SetDbObj(dbObj *sql.DB) error {
    db.dbObj = dbObj
    return nil
}

func getDb(ymlFile string) *sql.DB {
    data := getEnv(ymlFile)
    return dbConn(data)
}

func getEnv(yName string) map[interface{}]interface{} {
    ymlFile, err := ioutil.ReadFile(yName)

    if err != nil {
        log.Fatal(err)
    }

    data := make(map[interface{}]interface{})
    err2 := yaml.Unmarshal(ymlFile, &data)

    if err2 != nil {
        log.Fatal(err2)
    }
    return data
}

func dbConn(data map[interface{}]interface{}) *sql.DB {
    var connStr string = fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
        data["host"], data["port"], data["user"], data["password"], data["database"])

    dbObj, err := sql.Open("postgres", connStr)
    checkError(err)
    return dbObj
}

func (db *Pdb) Table() string {
    return db.table
}

func (db *Pdb) SetTable(table string) error {
    db.table = table
    return nil
}

func checkError(err error) {
    if err != nil {
        panic(err)
    }
}

func main() {
    pgDb := Pdb{}
    pgDb.SetDbObj(getDb("dbtest.yaml"))
    pgDb.SetTable("test_tbl")

    defer pgDb.dbObj.Close()
    dbPing(pgDb.dbObj)

    dropTable(pgDb)
    createTbl(pgDb)
    insertTbl(pgDb)

    rows := selectTbl(pgDb)
    defer rows.Close()

    printRows(rows)
    dropTable(pgDb)
}

func printRows(rows *sql.Rows) {
    var id int
    var name string
    var quantity int

    for rows.Next() {
        switch err := rows.Scan(&id, &name, &quantity); err {
        case sql.ErrNoRows:
            fmt.Println("No rows were returned")
        case nil:
            fmt.Printf("%d, %s, %d\n", id, name, quantity)
        default:
            checkError(err)
        }
    }
}

func selectTbl(pgDb Pdb) *sql.Rows {
    stmt1 := fmt.Sprintf("SELECT * from %s;", pgDb.table)
    rows, err := pgDb.dbObj.Query(stmt1)
    checkError(err)
    return rows
}

func insertTbl(pgDb Pdb) {
    stmtIns := fmt.Sprintf("INSERT INTO %s (name, quantity) VALUES ($1, $2);", pgDb.table)
    _, err := pgDb.dbObj.Exec(stmtIns, "test0", 100)
    checkError(err)
    _, err = pgDb.dbObj.Exec(stmtIns, "test1", 101)
    checkError(err)
    fmt.Println("Inserted 2 records")
}

func createTbl(pgDb Pdb) {
    creStr := fmt.Sprintf("CREATE TABLE %s (id serial PRIMARY KEY, name VARCHAR(20), quantity INTEGER);", pgDb.table)
    _, err := pgDb.dbObj.Exec(creStr)
    checkError(err)
    fmt.Println("Finished creating table")
}

func dropTable(pgDb Pdb) {
    dropStr := fmt.Sprintf("DROP TABLE IF EXISTS %s;", pgDb.table)
    _, err := pgDb.dbObj.Exec(dropStr)
    checkError(err)
    fmt.Println("Finished dropping table (if existed)")
}

func dbPing(db *sql.DB) {
    err := db.Ping()
    checkError(err)
    fmt.Println("Successfully created connection to database")
}