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


Драйвер 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.mod файл для отслеживания зависимостей go mod init кода 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.

Например, для хранилища 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

Подключение к 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". Необязательный 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"}),
  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. Дополнительные сведения об этом шаге см. в статье 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, используйте следующий фрагмент кода и пример кода в для подключения с использованием функции 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 можно использовать Azure CLI. См. получение токенов идентификатора Microsoft Entra для пользователей с помощью Azure CLI.

    Маркеры идентификатора 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, используйте следующий фрагмент кода и пример кода в для подключения с использованием функции 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"

Затем используйте следующий фрагмент кода и пример кода в Connect с помощью функции 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"

Затем используйте следующий фрагмент кода и пример кода в Connect с помощью функции 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 " + strconv.Itoa(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.

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