Работа с PostgreSQL в языке Go при помощи pgx

25 декабря 2019

Многие реальные приложения, написанные на Go, используют ту или иную РСУБД. Притом, последней нередко является PostgreSQL. Для работы с постгресом в мире Go существует больше одной библиотеки, в связи с чем возникает закономерный вопрос – а какую выбрать? Неплохим и достаточно популярным вариантом является jackc/pgx, с которым мы и познакомимся.

Модуль pgx имеет интерфейс, похожий на интерфейс database/sql, про который ранее уже рассказывалось в гостевом посте Владимира Солонина. Главное отличие между модулями заключается в том, что database/sql поддерживает разные РСУБД, а pgx – только PostgreSQL. За счет этого pgx может обеспечить большую производительность, и предложить больше возможностей, специфичных именно для постгреса. Если в приложении требуется поддерживать больше одной РСУБД одновременно (например, вы пишите CMS), берите database/sql. Иначе берите PostgreSQL и, соответственно, модуль pgx.

По устоявшейся традиции, писать будем телефонную книгу c REST-интерфейсом. Касаемо написания непосредственно REST на языке Go мне нечего добавить к другому гостевому посту Владимира, только вместо httprouter я решил использовать mux. Насколько мне известно, между ними нет особой разницы, только в mux хэндлеры имеют меньше аргументов. Субъективно такой код чуть красивее, только и всего. Для парсинга аргументов была использована cobra, а для парсинга файла конфигурации – viper.

На время разработки PostgreSQL удобно запускать через Docker:

# веб-страница образа: https://hub.docker.com/_/postgres
docker run -d --name postgresql -e POSTGRES_DB=restservice \
  -e POSTGRES_PASSWORD=s3cr3t -p 5432:5432 postgres:11

При необходимости получить доступ к базе данных можно так:

docker ps | grep postgresql
docker exec -it e7042bd737f8 psql -h localhost \
  -U postgres restservice -W

В приложении для подключения к СУБД будем использовать пул соединений, который уже реализован в pgx:

pool, err := pgxpool.Connect(context.Background(), dbURL)
if err != nil {
  log.Fatalf("Unable to connection to database: %v\n", err)
}
defer pool.Close()
log.Infof("Connected!")

Пример получения соединения из пула для выполнения миграции:

conn, err := pool.Acquire(context.Background())
if err != nil {
  log.Fatalf("Unable to acquire a database connection: %v\n", err)
}
migrateDatabase(conn.Conn())
conn.Release()

В самом pgx миграции не реализованы, но они есть в утилите jackc/tern. Там кода буквально на 300 строк. Они с легкостью переписываются на последнюю актуальную версию v4 модуля pgx, после чего в migrateDatabase остается написать только:

func migrateDatabase(conn *pgx.Conn) {
  migrator, err := migrate.NewMigrator(conn, "schema_version")
  if err != nil {
    log.Fatalf("Unable to create a migrator: %v\n", err)
  }

  err = migrator.LoadMigrations("./migrations")
  if err != nil {
    log.Fatalf("Unable to load migrations: %v\n", err)
  }

  err = migrator.Migrate()
  if err != nil {
    log.Fatalf("Unable to migrate: %v\n", err)
  }

  ver, err := migrator.GetCurrentVersion()
  if err != nil {
    log.Fatalf("Unable to get current schema version: %v\n", err)
  }

  log.Infof("Migration done. Current schema version: %v\n", ver)
}

Содержимое файла миграции:

CREATE TABLE phonebook(id SERIAL PRIMARY KEY,
                       name VARCHAR(64), phone VARCHAR(64));
---- create above / drop below ----
DROP TABLE phonebook;

Таким образом, приложение может само создавать себе схему базы данных или изменять ее при обновлении. Локи реализованы на функции pg_advisory_lock, что исключает возможность запуска двух миграций параллельно.

Далее все очень похоже на database/sql. Так, например, происходит вставка новой записи:

type Record struct {
  Id    int    `json:"id"`
  Name  string `json:"name"`
  Phone string `json:"phone"`
}

func Insert(p *pgxpool.Pool, w http.ResponseWriter, r *http.Request) {
  var rec Record
  err := json.NewDecoder(r.Body).Decode(&rec)
  if err != nil { // bad request
    w.WriteHeader(400)
    return
  }

  conn, err := p.Acquire(context.Background())
  if err != nil {
    log.Errorf("Unable to acquire a database connection: %v\n", err)
    w.WriteHeader(500)
    return
  }
  defer conn.Release()

  row := conn.QueryRow(context.Background(),
    "INSERT INTO phonebook (name, phone) VALUES ($1, $2) RETURNING id",
    rec.Name, rec.Phone)
  var id uint64
  err = row.Scan(&id)
  if err != nil {
    log.Errorf("Unable to INSERT: %v\n", err)
    w.WriteHeader(500)
    return
  }

  resp := make(map[string]string, 1)
  resp["id"] = strconv.FormatUint(id, 10)
  w.Header().Set("Content-Type", "application/json; charset=utf-8")
  err = json.NewEncoder(w).Encode(resp)
  if err != nil {
    log.Errorf("Unable to encode json: %v\n", err)
    w.WriteHeader(500)
    return
  }
}

А так – удаление:

func Delete(p *pgxpool.Pool, w http.ResponseWriter, r *http.Request) {
  vars := mux.Vars(r)
  id, err := strconv.ParseUint(vars["id"], 10, 64)
  if err != nil { // bad request
    w.WriteHeader(400)
    return
  }

  conn, err := p.Acquire(context.Background())
  if err != nil {
    log.Errorf("Unable to acquire a database connection: %v\n", err)
    w.WriteHeader(500)
    return
  }
  defer conn.Release()

  ct, err := conn.Exec(context.Background(),
    "DELETE FROM phonebook WHERE id = $1", id)
  if err != nil {
    log.Errorf("Unable to DELETE: %v\n", err)
    w.WriteHeader(500)
    return
  }

  if ct.RowsAffected() == 0 {
    w.WriteHeader(404)
    return
  }

  w.WriteHeader(200)
}

Полную версию исходников вы найдете в этом репозитории на GitHub. Примеры запросов к сервису:

# создание:
curl -vvv -XPOST -H 'Content-Type: application/json' \
  -d '{"name":"Alice","phone":"123"}'  localhost:8080/api/v1/records

# чтение:
curl -vvv localhost:8080/api/v1/records/123

# обновление:
curl -vvv -XPUT -H 'Content-Type: application/json' \
  -d '{"name":"Bob","phone":"456"}' localhost:8080/api/v1/records/456

# удаление:
curl -vvv -XDELETE  localhost:8080/api/v1/records/789

Заинтересованные читатели могут подумать над тем, как написать тесты к этому коду, а также реализовать API GET /api/v1/records, который сейчас показывает лишь сообщение «Under construction». API возвращает все записи, которые имеются в базе, но не более 1000 штук за один запрос. Принимаемые аргументы: limit, offset и order. Первый аргумент опционален и позволяет возвращать менее 1000 записей. Через второй аргумент можно передать id последней записи, которую вернул предыдущий запрос. Тогда текущий запрос вернет записи, идущие после offset. Наконец, order – это строковый параметр, который может быть либо asc (сортировка по возрастанию id), либо desc (по убыванию). Если order не указан, считается, что он asc.

Дополнение: Вас также могут заинтересовать посты Генерация SQL-запросов в Go с помощью squirrel и Тестирование проектов на Go с dockertest.

Дополнение: Я переписал код так, чтобы сервис был совместим с CockroachDB.