Поделиться через

Драйвер SQL Databricks для Go

Драйвер SQL Databricks для Go — это библиотека Go, которая позволяет использовать код Go для выполнения команд SQL в вычислительных ресурсах Azure Databricks. В этой статье приведены примеры драйвера SQL Databricks для Go README, справочник по API и примеры.

Требования

Начало работы с драйвером SQL Databricks для Go

  1. На вашем компьютере для разработки, на котором уже установлена версия Go 1.20 или более поздняя, и уже создан существующий проект на Go, создайте файл для отслеживания зависимостей вашего кода Go, выполнив команду , например,

    go mod init sample

  2. Воспользуйтесь пакетом Databricks SQL Driver for Go, выполнив команду go mod edit -require, заменив v1.5.2 на последнюю версию пакета, указанную в Releases.

    go mod edit -require github.com/databricks/databricks-sql-go@v1.5.2

    Теперь файл go.mod должен выглядеть так:

    module sample

go 1.20

require github.com/databricks/databricks-sql-go v1.5.2

  3. В проекте создайте файл кода Go, который импортирует драйвер SQL Databricks для Go. В следующем примере в файле main.go со следующим содержимым перечислены все кластеры в рабочей области Azure Databricks:

    package main

import (
  "database/sql"
  "os"
  _ "github.com/databricks/databricks-sql-go"
)

func main() {
  dsn := os.Getenv("DATABRICKS_DSN")

  if dsn == "" {
    panic("No connection string found. " +
     "Set the DATABRICKS_DSN environment variable, and try again.")
  }

  db, err := sql.Open("databricks", dsn)
  if err != nil {
    panic(err)
  }
  defer db.Close()

  if err := db.Ping(); err != nil {
    panic(err)
  }
}

  4. Добавьте все отсутствующие зависимости модуля, выполнив go mod tidy команду:

    go mod tidy

    Примечание.

    Если вы получаете ошибку go: warning: "all" matched no packages, вы забыли добавить файл кода Go, который импортирует драйвер SQL Databricks для Go.

  5. Создайте копии всех пакетов, необходимых для поддержки сборок и тестов пакетов в main модуле go mod vendor , выполнив команду:

    go mod vendor

  6. Измените код по мере необходимости, чтобы задать переменную среды DATABRICKS_DSN для azure Databricks аутентификации. См. также подключение с помощью строки подключения DSN.

  7. Запустите файл кода Go, предполагая, что файл называется main.go, выполнив команду go run.

    go run main.go

  8. Если ошибки не возвращаются, вы успешно выполнили проверку подлинности драйвера SQL Databricks для Go с рабочей областью Azure Databricks и подключились к работающему кластеру Azure Databricks или SQL хранилищу в этой рабочей области.

Подключение по строке подключения DSN

Чтобы получить доступ к кластерам и хранилищам SQL, используйте sql.Open() для создания дескриптора базы данных через DSN-строку подключения. В этом примере кода извлекается строка подключения DSN из переменной среды с именемDATABRICKS_DSN:

package main

import (
  "database/sql"
  "os"
  _ "github.com/databricks/databricks-sql-go"
)

func main() {
  dsn := os.Getenv("DATABRICKS_DSN")

  if dsn == "" {
    panic("No connection string found. " +
          "Set the DATABRICKS_DSN environment variable, and try again.")
  }

  db, err := sql.Open("databricks", dsn)
  if err != nil {
    panic(err)
  }
  defer db.Close()

  if err := db.Ping(); err != nil {
    panic(err)
  }
}

Чтобы указать строку подключения DSN в правильном формате, см. примеры строки подключения DSN в разделе Аутентификация. Например, для аутентификации с использованием личного токена доступа Azure Databricks используйте следующий синтаксис, где:

  • <personal-access-token> представляет личный токен доступа Azure Databricks согласно требованиям.
  • <server-hostname> — это значение имени узла сервера из раздела "Требования";
  • <port-number> — это значение порта из требований, которое обычно 443.
  • <http-path> — это значение пути HTTP из раздела "Требования".
  • <paramX=valueX> — это один или несколько необязательных параметров, перечисленных далее в этой статье.
token:<personal-access-token>@<server-hostname>:<port-number>/<http-path>?<param1=value1>&<param2=value2>

Например, для кластера:

token:dapi12345678901234567890123456789012@adb-1234567890123456.7.azuredatabricks.net:443/sql/protocolv1/o/1234567890123456/1234-567890-abcdefgh

Например, для хранилища SQL:

token:dapi12345678901234567890123456789012@adb-1234567890123456.7.azuredatabricks.net:443/sql/1.0/endpoints/a1b234c5678901d2

Примечание.

Чтобы обеспечить безопасность, не рекомендуется жестко кодировать эту строку подключения DSN в код Go. Вместо этого следует получить эту строку подключения DSN из безопасного расположения. Например, пример кода, приведенный ранее в этой статье, использовал переменную среды.

Необязательные параметры

  • Поддерживаемые необязательные параметры подключения можно указать в <param=value>. Некоторые из наиболее часто используемых включают:
    • catalog: задает имя начального каталога в сеансе.
    • schema: задает имя начальной схемы в сеансе.
    • maxRows: настраивает максимальное количество строк, извлекаемых на один запрос. Значение по умолчанию — 10000.
    • timeout: добавляет время ожидания (в секундах) для выполнения запроса сервера. Значение по умолчанию — отсутствие таймаута.
    • userAgentEntry: используется для идентификации партнеров. Дополнительные сведения см. в документации партнера.
  • Поддерживаемые необязательные параметры сеанса можно указать в param=value. Некоторые из наиболее часто используемых включают:
    • ansi_mode: Булевая строка. true для инструкций сеанса, которые соответствуют правилам, заданным спецификацией ANSI SQL. По умолчанию система имеет значение false.
    • timezone: строка, например America/Los_Angeles. Задает часовой пояс сеанса. Системное значение по умолчанию — UTC.
    • query_tags: строка пар "ключ:значение", разделенных запятыми, для добавления к запросам SQL для отслеживания в system.query.history, например team:engineering,dashboard:abc123. Эта функция доступна в общедоступной предварительной версии.

Например, для хранилища SQL:

token:dapi12345678901234567890123456789012@adb-1234567890123456.7.azuredatabricks.net:443/sql/1.0/endpoints/a1b234c5678901d2?catalog=hive_metastore&schema=example&maxRows=100&timeout=60&timezone=America/Sao_Paulo&ansi_mode=true&query_tags=team:engineering,env:prod

Подключение к NewConnector функции

Кроме того, используйте sql.OpenDB() для создания дескриптора базы данных с новым объектом соединителя, созданным с помощью dbsql.NewConnector(). Для подключения к кластерам Azure Databricks и хранилищам SQL с новым объектом соединителя требуется версия 1.0.0 или более поздняя драйвера SQL Databricks для Go. Например:

package main

import (
  "database/sql"
  "os"
  dbsql "github.com/databricks/databricks-sql-go"
)

func main() {
  connector, err := dbsql.NewConnector(
    dbsql.WithAccessToken(os.Getenv("DATABRICKS_ACCESS_TOKEN")),
    dbsql.WithServerHostname(os.Getenv("DATABRICKS_HOST")),
    dbsql.WithPort(443),
    dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTP_PATH")),
  )
  if err != nil {
    panic(err)
  }

  db := sql.OpenDB(connector)
  defer db.Close()

  if err := db.Ping(); err != nil {
    panic(err)
  }
}

Чтобы указать правильный набор параметров NewConnector, см. примеры в аутентификации.

Примечание.

В качестве рекомендации по обеспечению безопасности не следует жестко кодировать NewConnector параметры в код Go. Вместо этого следует получить эти значения из безопасного расположения. Например, предыдущий код использует переменные среды.

Ниже перечислены некоторые из наиболее часто используемых функциональных параметров:

  • WithAccessToken(<access-token>): ваш личный токен доступа Azure Databricks из раздела требований. Обязательный string.
  • WithServerHostname(<server-hostname>)Значение имени узла сервера из требований. Обязательный string.
  • WithPort(<port>): номер порта сервера, как правило 443. Обязательный int.
  • WithHTTPPath(<http-path>) Значение HTTP Path из требований. Обязательный string.
  • WithInitialNamespace(<catalog>, <schema>):имя каталога и схемы в сеансе. Необязательный string, string.
  • WithMaxRows(<max-rows>): максимальное количество строк, извлекаемых в запросе. Значение по умолчанию — 10000. Необязательно int.
  • WithSessionParams(<params-map>): параметры сеанса, включая "часовой пояс", "ansi_mode" и "query_tags". Необязательный map[string]string.
  • WithTimeout(<timeout>). Время ожидания (в time.Duration) для выполнения запроса сервера. Значение по умолчанию — отсутствие таймаута. Необязательно.
  • WithUserAgentEntry(<isv-name-plus-product-name>). Используется для идентификации партнеров. Дополнительные сведения см. в документации партнера. Необязательный string.

Например:

connector, err := dbsql.NewConnector(
  dbsql.WithAccessToken(os.Getenv("DATABRICKS_ACCESS_TOKEN")),
  dbsql.WithServerHostname(os.Getenv("DATABRICKS_HOST")),
  dbsql.WithPort(443),
  dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTP_PATH")),
  dbsql.WithInitialNamespace("samples", "nyctaxi"),
  dbsql.WithMaxRows(100),
  dbsql.SessionParams(map[string]string{"timezone": "America/Sao_Paulo", "ansi_mode": "true", "query_tags": "team:analytics,project:reporting"}),
  dbsql.WithTimeout(time.Minute),
  dbsql.WithUserAgentEntry("example-user"),
)

Проверка подлинности

Драйвер SQL Databricks для Go поддерживает следующие типы проверки подлинности Azure Databricks:

Драйвер SQL Databricks для Go пока не поддерживает следующие типы проверки подлинности Azure Databricks:

Проверка подлинности маркера личного доступа Databricks

Чтобы использовать SQL-драйвер Databricks для Go с аутентификацией токена личного доступа Azure Databricks , необходимо сначала создать личный токен доступа Azure Databricks. Дополнительные сведения об этом шаге см. в разделе "Создание личных маркеров доступа для пользователей рабочей области".

Чтобы выполнить проверку подлинности "Databricks SQL Driver for Go" с использованием строки подключения DSN и примера кода в Connect с подключением через строку DSN, используйте следующий синтаксис строки подключения DSN, где:

  • <personal-access-token> представляет личный токен доступа Azure Databricks согласно требованиям.
  • <server-hostname> — это значение имени узла сервера из раздела "Требования";
  • <port-number> — это значение порта из требований, которое обычно 443.
  • <http-path> — это значение пути HTTP из раздела "Требования".

Вы также можете добавить один или несколько необязательных параметров, перечисленных ранее в этой статье.

token:<personal-access-token>@<server-hostname>:<port-number>/<http-path>

Чтобы аутентифицировать Databricks SQL Driver for Go с помощью функции NewConnector, используйте следующий пример кода и пример кода в Connect с функцией NewConnector, который предполагает, что вы установили следующие переменные среды.

  • Установите на значение имени узла сервера для вашего кластера или SQL-склада.
  • Установите для DATABRICKS_HTTP_PATHзначение пути HTTP для вашего кластера или SQL-сервера.
  • Назначьте DATABRICKS_TOKENдля маркера личного доступа Azure Databricks.

Чтобы задать переменные среды, ознакомьтесь с документацией операционной системы.

connector, err := dbsql.NewConnector(
  dbsql.WithServerHostname(os.Getenv("DATABRICKS_SERVER_HOSTNAME")),
  dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTP_PATH")),
  dbsql.WithPort(443),
  dbsql.WithAccessToken(os.Getenv("DATABRICKS_TOKEN")),
)

Проверка подлинности токена Microsoft Entra ID

Драйвер Databricks SQL для Go поддерживает токены Microsoft Entra ID для пользователя Azure Databricks или сервисного субъекта Microsoft Entra ID.

Чтобы создать токен доступа Entra ID Microsoft, сделайте следующее:

  • Для пользователя Azure Databricks или учетной записи службы Microsoft Entra ID используйте Azure CLI. См. Получение токенов Microsoft Entra ID вручную. Чтобы создать управляемый принципал службы Microsoft Entra ID, см. раздел Субъекты-службы.

    Токены идентификатора Microsoft Entra имеют продолжительность действия по умолчанию около 1 часа. Чтобы создать новый маркер идентификатора Microsoft Entra, повторите этот процесс.

    Чтобы выполнить проверку подлинности "Databricks SQL Driver for Go" с использованием строки подключения DSN и примера кода в Connect с подключением через строку DSN, используйте следующий синтаксис строки подключения DSN, где:

    • <microsoft-entra-id-token> — это маркер идентификатора Microsoft Entra.
    • <server-hostname> — это значение имени узла сервера из раздела "Требования";
    • <port-number> — это значение порта из требований, которое обычно 443.
    • <http-path> — это значение пути HTTP из раздела "Требования".

    Вы также можете добавить один или несколько необязательных параметров, перечисленных ранее в этой статье.

    token:<microsoft-entra-id-token>@<server-hostname>:<port-number>/<http-path>

    Чтобы аутентифицировать Databricks SQL Driver for Go с помощью функции NewConnector, используйте следующий пример кода и пример кода в Connect с функцией NewConnector, который предполагает, что вы установили следующие переменные среды.

    • Установите на значение имени узла сервера для вашего кластера или SQL-склада.
    • Установите для DATABRICKS_HTTP_PATHзначение пути HTTP для вашего кластера или SQL-сервера.
    • Установите DATABRICKS_TOKENдля вашего маркера идентификатора Microsoft Entra.

    Чтобы задать переменные среды, ознакомьтесь с документацией операционной системы.

    connector, err := dbsql.NewConnector(
  dbsql.WithServerHostname(os.Getenv("DATABRICKS_SERVER_HOSTNAME")),
  dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTP_PATH")),
  dbsql.WithPort(443),
  dbsql.WithAccessToken(os.Getenv("DATABRICKS_TOKEN")),
)

Аутентификация OAuth U2M от пользователя к машине

Databricks SQL Driver for Go версии 1.5.0 и более поздних версий поддерживают проверку подлинности OAuth на компьютере (U2M).

Чтобы использовать Драйвер SQL Databricks для Go со строкой подключения DSN и примером кода в Connect с строкой подключения DSN, используйте следующий синтаксис строки подключения DSN, где:

  • <server-hostname> — это значение имени узла сервера из раздела "Требования";
  • <port-number> — это значение порта из требований, которое обычно 443.
  • <http-path> — это значение пути HTTP из раздела "Требования".

Вы также можете добавить один или несколько необязательных параметров, перечисленных ранее в этой статье.

<server-hostname>:<port-number>/<http-path>?authType=OauthU2M

Чтобы проверить подлинность драйвера SQL Databricks для Go с помощью функции NewConnector, необходимо сначала добавить следующее в объявление import.

"github.com/databricks/databricks-sql-go/auth/oauth/u2m"

Затем используйте следующий фрагмент кода и пример в соединении с функцией NewConnector, который предполагает, что вы задали следующие переменные окружения:

  • Установите на значение имени узла сервера для вашего кластера или SQL-склада.
  • Установите для DATABRICKS_HTTP_PATHзначение пути HTTP для вашего кластера или SQL-сервера.

Чтобы задать переменные среды, ознакомьтесь с документацией операционной системы.

authenticator, err := u2m.NewAuthenticator(os.Getenv("DATABRICKS_SERVER_HOSTNAME"), 1*time.Minute)
if err != nil {
  panic(err)
}

connector, err := dbsql.NewConnector(
  dbsql.WithServerHostname(os.Getenv("DATABRICKS_SERVER_HOSTNAME")),
  dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTP_PATH")),
  dbsql.WithPort(443),
  dbsql.WithAuthenticator(authenticator),
)

Аутентификация OAuth между машинами (M2M)

Databricks SQL Driver for Go версии 1.5.2 и выше поддерживают проверку подлинности OAuth на компьютере (M2M).

Чтобы использовать драйвер SQL Databricks для Go с проверкой подлинности OAuth M2M, необходимо выполнить следующее:

  1. Создайте субъект-службу Azure Databricks в рабочей области Azure Databricks и создайте секрет OAuth для этого субъекта-службы.

    Сведения о создании сервисного субъекта и его секрета OAuth см. в разделе "Авторизация доступа сервисного субъекта к Azure Databricks через OAuth". Запишите значение UUID или идентификатора приложения принципала службы и значение секрета для OAuth-ключа принципала службы.

  2. Предоставьте субъекту-службе доступ к кластеру или хранилищу.

    Чтобы предоставить служебному принципалу доступ к кластеру или хранилищу, см. разрешения на вычислительные ресурсы или управление SQL-хранилищем.

Чтобы выполнить проверку подлинности "Databricks SQL Driver for Go" с использованием строки подключения DSN и примера кода в Connect с подключением через строку DSN, используйте следующий синтаксис строки подключения DSN, где:

  • <server-hostname> — это значение имени узла сервера из раздела "Требования";
  • <port-number> — это значение порта из требований, которое обычно 443.
  • <http-path> — это значение пути HTTP из раздела "Требования".
  • <client-id> — это значение UUID субъекта-службы или идентификатора приложения .
  • <client-secret> — это секретное значение для OAuth ключа субъекта-службы.

Вы также можете добавить один или несколько необязательных параметров, перечисленных ранее в этой статье.

<server-hostname>:<port-number>/<http-path>?authType=OAuthM2M&clientID=<client-id>&clientSecret=<client-secret>

Чтобы проверить подлинность драйвера SQL Databricks для Go с функцией NewConnector, сначала необходимо добавить следующее к объявлению import.

"github.com/databricks/databricks-sql-go/auth/oauth/m2m"

Затем используйте следующий фрагмент кода и пример в соединении с функцией NewConnector, который предполагает, что вы задали следующие переменные окружения:

  • Установите на значение имени узла сервера для вашего кластера или SQL-склада.
  • Установите для DATABRICKS_HTTP_PATHзначение пути HTTP для вашего кластера или SQL-сервера.
  • DATABRICKS_CLIENT_IDУстановите для главного объекта службы значение UUID или идентификатор приложения.
  • DATABRICKS_CLIENT_SECRET, установите значение секрет для OAuth секрета служебного принципала.

Чтобы задать переменные среды, ознакомьтесь с документацией операционной системы.

authenticator := m2m.NewAuthenticator(
  os.Getenv("DATABRICKS_CLIENT_ID"),
  os.Getenv("DATABRICKS_CLIENT_SECRET"),
  os.Getenv("DATABRICKS_SERVER_HOSTNAME"),
)

connector, err := dbsql.NewConnector(
  dbsql.WithServerHostname(os.Getenv("DATABRICKS_SERVER_HOSTNAME")),
  dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTP_PATH")),
  dbsql.WithPort(443),
  dbsql.WithAuthenticator(authenticator),
)

Запрос данных

В следующем примере кода демонстрируется вызов драйвера Databricks SQL для Go с целью выполнения простого SQL-запроса к вычислительному ресурсу Azure Databricks. Эта команда возвращает первые две строки из таблицы trips в схеме samples каталога nyctaxi.

Приведенный ниже код извлекает строку подключения DSN из переменной среды с именем DATABRICKS_DSN.

package main

import (
  "database/sql"
  "fmt"
  "os"
  "time"

  _ "github.com/databricks/databricks-sql-go"
)

func main() {
  dsn := os.Getenv("DATABRICKS_DSN")

  if dsn == "" {
    panic("No connection string found." +
          "Set the DATABRICKS_DSN environment variable, and try again.")
  }

  db, err := sql.Open("databricks", dsn)
  if err != nil {
    panic(err)
  }

  defer db.Close()

  var (
    tpep_pickup_datetime  time.Time
    tpep_dropoff_datetime time.Time
    trip_distance         float64
    fare_amount           float64
    pickup_zip            int
    dropoff_zip           int
  )

  rows, err := db.Query("SELECT * FROM samples.nyctaxi.trips LIMIT ?", 2)
  if err != nil {
    panic(err)
  }

  defer rows.Close()

  fmt.Print("tpep_pickup_datetime,",
    "tpep_dropoff_datetime,",
    "trip_distance,",
    "fare_amount,",
    "pickup_zip,",
    "dropoff_zip\n")

  for rows.Next() {
    err := rows.Scan(&tpep_pickup_datetime,
      &tpep_dropoff_datetime,
      &trip_distance,
      &fare_amount,
      &pickup_zip,
      &dropoff_zip)
    if err != nil {
      panic(err)
    }

    fmt.Print(tpep_pickup_datetime, ",",
      tpep_dropoff_datetime, ",",
      trip_distance, ",",
      fare_amount, ",",
      pickup_zip, ",",
      dropoff_zip, "\n")
  }

  err = rows.Err()
  if err != nil {
    panic(err)
  }
}

Управление файлами в томах каталога Unity

Драйвер SQL Databricks позволяет записывать локальные файлы в каталог Unity томов, загружать файлы из томов и удалять файлы из томов, как показано в следующем примере:

package main

import (
  "context"
  "database/sql"
  "os"

  _ "github.com/databricks/databricks-sql-go"
  "github.com/databricks/databricks-sql-go/driverctx"
)

func main() {
  dsn := os.Getenv("DATABRICKS_DSN")

  if dsn == "" {
    panic("No connection string found." +
      "Set the DATABRICKS_DSN environment variable, and try again.")
  }

  db, err := sql.Open("databricks", dsn)
  if err != nil {
    panic(err)
  }
  defer db.Close()

  // For writing local files to volumes and downloading files from volumes,
  // you must first specify the path to the local folder that contains the
  // files to be written or downloaded.
  // For multiple folders, add their paths to the following string array.
  // For deleting files in volumes, this string array is ignored but must
  // still be provided, so in that case its value can be set for example
  // to an empty string.
  ctx := driverctx.NewContextWithStagingInfo(
    context.Background(),
    []string{"/tmp/"},
  )

  // Write a local file to the path in the specified volume.
  // Specify OVERWRITE to overwrite any existing file in that path.
  db.ExecContext(ctx, "PUT '/tmp/my-data.csv' INTO '/Volumes/main/default/my-volume/my-data.csv' OVERWRITE")

  // Download a file from the path in the specified volume.
  db.ExecContext(ctx, "GET '/Volumes/main/default/my-volume/my-data.csv' TO '/tmp/my-downloaded-data.csv'")

  // Delete a file from the path in the specified volume.
  db.ExecContext(ctx, "REMOVE '/Volumes/main/default/my-volume/my-data.csv'")

  db.Close()
}

Ведение журнала

Используется github.com/databricks/databricks-sql-go/logger для регистрации сообщений, которые выдает драйвер SQL Databricks для Go. В следующем примере кода используется sql.Open() для создания дескриптора базы данных с помощью строки подключения DSN. В этом примере кода извлекается строка подключения DSN из переменной среды с именемDATABRICKS_DSN. Все сообщения журнала, создаваемые на debug уровне и ниже, записываются в файл results.log.

package main

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

  _ "github.com/databricks/databricks-sql-go"
  dbsqllog "github.com/databricks/databricks-sql-go/logger"
)

func main() {
  dsn := os.Getenv("DATABRICKS_DSN")

  // Use the specified file for logging messages to.
  file, err := os.Create("results.log")
  if err != nil {
    log.Fatal(err)
  }
  defer file.Close()

  writer := io.Writer(file)

  // Log messages at the debug level and below.
  if err := dbsqllog.SetLogLevel("debug"); err != nil {
    log.Fatal(err)
  }

  // Log messages to the file.
  dbsqllog.SetLogOutput(writer)

  if dsn == "" {
    panic("Error: Cannot connect. No connection string found. " +
      "Set the DATABRICKS_DSN environment variable, and try again.")
  }

  db, err := sql.Open("databricks", dsn)
  if err != nil {
    panic(err)
  }
  defer db.Close()

  if err := db.Ping(); err != nil {
    panic(err)
  }
}

Тестирование

Чтобы протестировать код, используйте платформы тестирования Go, такие как стандартная библиотека тестирования . Чтобы протестировать код в имитированных условиях, не вызывая конечные точки Azure Databricks REST API или не изменяя состояние ваших Azure Databricks учетных записей или рабочих пространств, используйте библиотеки имитации Go, такие как testfify.

Например, если у вас есть следующий файл с именем helpers.go, он содержит функцию GetDBWithDSNPAT, которая возвращает подключение к рабочей области Azure Databricks; функцию GetNYCTaxiTrips, которая извлекает данные из таблицы trips в схеме samples каталога nyctaxi; и функцию PrintNYCTaxiTrips, которая выводит полученные данные.

package main

import (
  "database/sql"
  "fmt"
  "strconv"
  "time"
)

func GetDBWithDSNPAT(dsn string) (*sql.DB, error) {
  db, err := sql.Open("databricks", dsn)
  if err != nil {
    return nil, err
  }
  return db, nil
}

func GetNYCTaxiTrips(db *sql.DB, numRows int) (*sql.Rows, error) {
  rows, err := db.Query("SELECT * FROM samples.nyctaxi.trips LIMIT ?", numRows)
  if err != nil {
    return nil, err
  }
  return rows, nil
}

func PrintNYCTaxiTrips(rows *sql.Rows) {
  var (
    tpep_pickup_datetime  time.Time
    tpep_dropoff_datetime time.Time
    trip_distance         float64
    fare_amount           float64
    pickup_zip            int
    dropoff_zip           int
  )

  fmt.Print(
    "tpep_pickup_datetime,",
    "tpep_dropoff_datetime,",
    "trip_distance,",
    "fare_amount,",
    "pickup_zip,",
    "dropoff_zip\n",
  )

  for rows.Next() {
    err := rows.Scan(
      &tpep_pickup_datetime,
      &tpep_dropoff_datetime,
      &trip_distance,
      &fare_amount,
      &pickup_zip,
      &dropoff_zip,
    )
    if err != nil {
      panic(err)
    }

    fmt.Print(
      tpep_pickup_datetime, ",",
      tpep_dropoff_datetime, ",",
      trip_distance, ",",
      fare_amount, ",",
      pickup_zip, ",",
      dropoff_zip, "\n",
    )
  }

  err := rows.Err()
  if err != nil {
    panic(err)
  }
}

Учитывая указанный файл с именем main.go, который вызывает эти функции:

package main

import (
  "os"
)

func main() {
  db, err := GetDBWithDSNPAT(os.Getenv("DATABRICKS_DSN"))
  if err != nil {
    panic(err)
  }

  rows, err := GetNYCTaxiTrips(db, 2)
  if err != nil {
    panic(err)
  }

  PrintNYCTaxiTrips(rows)
}

Следующий файл с именем helpers_test.go проверяет, возвращает ли GetNYCTaxiTrips функция ожидаемый ответ. Вместо создания реального подключения к целевой рабочей области этот тест макетирует sql.DB объект. Тест также макетирует некоторые данные, соответствующие схеме и значениям, которые находятся в реальных данных. Тест возвращает мокаемые данные через мокаемое подключение, а затем проверяет, соответствует ли одно из значений строк мокаемых данных ожидаемому значению.

package main

import (
  "database/sql"
  "testing"

  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/mock"
)

// Define an interface that contains a method with the same signature
// as the real GetNYCTaxiTrips function that you want to test.
type MockGetNYCTaxiTrips interface {
  GetNYCTaxiTrips(db *sql.DB, numRows int) (*sql.Rows, error)
}

// Define a struct that represents the receiver of the interface's method
// that you want to test.
type MockGetNYCTaxiTripsObj struct {
  mock.Mock
}

// Define the behavior of the interface's method that you want to test.
func (m *MockGetNYCTaxiTripsObj) GetNYCTaxiTrips(db *sql.DB, numRows int) (*sql.Rows, error) {
  args := m.Called(db, numRows)
  return args.Get(0).(*sql.Rows), args.Error(1)
}

func TestGetNYCTaxiTrips(t *testing.T) {
  // Instantiate the receiver.
  mockGetNYCTaxiTripsObj := new(MockGetNYCTaxiTripsObj)

  // Define how the mock function should be called and what it should return.
  // We're not concerned with whether the actual database is connected to--just
  // what is returned.
  mockGetNYCTaxiTripsObj.On("GetNYCTaxiTrips", mock.Anything, mock.AnythingOfType("int")).Return(&sql.Rows{}, nil)

  // Call the mock function that you want to test.
  rows, err := mockGetNYCTaxiTripsObj.GetNYCTaxiTrips(nil, 2)

  // Assert that the mock function was called as expected.
  mockGetNYCTaxiTripsObj.AssertExpectations(t)

  // Assert that the mock function returned what you expected.
  assert.NotNil(t, rows)
  assert.Nil(t, err)
}

Так как функция GetNYCTaxiTrips содержит инструкцию SELECT и поэтому не изменяет состояние таблицы trips, макетирование не является абсолютно обязательным в этом примере. Однако имитация позволяет быстро проводить тесты без ожидания реального подключения к рабочей области. Кроме того, макетирование позволяет выполнять имитированные тесты несколько раз для функций, которые могут изменить состояние таблицы, например INSERT INTO, UPDATEи DELETE FROM.

Дополнительные ресурсы

  • Last updated on