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

Драйвер 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 последнюю версию пакета Databricks SQL Driver for v1.5.2 Go, как указано в выпусках:

    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 из требований. Обязательный 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

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

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

  • Для пользователя 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),
)

Проверка подлинности на компьютере (M2M) OAuth

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, такие как стандартная библиотека тестирования . Чтобы протестировать код в имитированных условиях, не вызывая конечные точки REST API Azure Databricks или изменяя состояние учетных записей Или рабочих областей Azure Databricks, используйте библиотеки Go, такие как проверка.

Например, если у вас есть следующий файл с именем 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