Conector do Databricks SQL para Python

O Conector do Databricks SQL para Python é uma biblioteca Python que permite a você usar o código Python para executar comandos SQL em clusters do Azure Databricks e em warehouses do Databricks SQL. O Conector SQL Databricks para Python é mais fácil de configurar e usar do que bibliotecas do Python semelhantes, como pyodbc. Essa biblioteca segue PEP 249 – Especificação da API do Banco de Dados Python v2.0.

Requisitos

• Um computador de desenvolvimento executando Python >=3.8 e <=3.11.
• O Databricks recomenda que você use ambientes virtuais Python, como os fornecidos por venv que estão incluídos no Python. Os ambientes virtuais ajudam a garantir que você esteja usando as versões corretas do Python e do Conector SQL do Databricks para Python. A configuração e o uso de ambientes virtuais estão fora do escopo deste artigo. Para obter mais informações, confira Criando ambientes virtuais.
• Um cluster ou SQL warehouse existente.

Começar

• Instale o Conector SQL do Databricks para o Python. PyArrow é uma dependência opcional do Databricks SQL Connector para Python e não é instalado por padrão. Se você não instalar o PyArrow, recursos como o CloudFetch e outras funcionalidades do Apache Arrow não estarão disponíveis, o que pode afetar o desempenho de grandes volumes de dados.

  • Para instalar o conector lean, use pip install databricks-sql-connector.
  • Para instalar o conector completo, incluindo o PyArrow, use pip install databricks-sql-connector[pyarrow].

• Reúna as seguintes informações sobre o cluster ou SQL warehouse que você deseja usar:

  Agrupamento

  • O nome do host do servidor do cluster. Você pode obter isso do valor de Nome do Host do Servidor na guia Opções Avançadas > JDBC/ODBC do cluster.
  • O caminho HTTP do cluster. Você pode obter isso do valor de caminho HTTP do Servidor na guia Opções Avançadas > JDBC/ODBC do cluster.

  SQL warehouse

  • O nome do host do servidor do SQL warehouse. É possível obtê-lo no valor do Nome do host do servidor na guia Detalhes da conexão do SQL warehouse.
  • O caminho HTTP do SQL warehouse. É possível obtê-lo no valor de Caminho HTTP na guia Detalhes da conexão do SQL warehouse.

Autenticação

O Conector do SQL do Databricks para Python dá suporte para os seguintes tipos de autenticação do Azure Databricks:

• Autenticação de token de acesso pessoal do Databricks
• Autenticação de token do Microsoft Entra ID
• Autenticação OAuth entre máquinas (M2M)
• Autenticação U2M (usuário para computador) do OAuth

O Conector SQL do Databricks para Python ainda não suporta os seguintes tipos de autenticação do Azure Databricks:

• Autenticação de identidades gerenciadas do Azure
• Autenticação da entidade de serviço do MS Entra
• Autenticação da CLI do Azure

Autenticação de token de acesso pessoal do Databricks

Para usar o Conector do SQL do Databricks para Python com a autenticação de token de acesso pessoal do Azure Databricks, você precisa primeiro criar um token de acesso pessoal do Azure Databricks. Para fazer isso, siga as etapas em Tokens de acesso pessoal do Azure Databricks para usuários do workspace.

Para autenticar o Databricks SQL Connector for Python, use o seguinte trecho de código. Esse snippet de código pressupõe que você tenha definido as seguintes variáveis de ambiente:

• DATABRICKS_SERVER_HOSTNAMEdefinido como o valor Nome do Host do Servidor do seu cluster ou SQL Warehouse.
• DATABRICKS_HTTP_PATH, definido para o valor do Caminho HTTP para seu cluster ou SQL Warehouse.
• DATABRICKS_TOKEN, definido como o token de acesso pessoal do Azure Databricks.

Para definir variáveis de ambiente, confira a documentação do sistema operacional.

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:
# ...

Autenticação M2M (de computador para computador) do OAuth

O Conector do Databricks SQL para Python versões 2.7.0 e superior suportam autenticação OAuth máquina a máquina (M2M). Você também precisa instalar o SDK do Databricks para Python 0.18.0 ou superior, por exemplo, executando pip install databricks-sdk ou python -m pip install databricks-sdk.

Para usar o Conector do Databricks SQL para Python com autenticação OAuth M2M, faça o seguinte:

1. Crie um principal de serviço do Azure Databricks em seu espaço de trabalho do Azure Databricks e crie um segredo OAuth para esse principal de serviço.

   Para criar a entidade de serviço e o respectivo segredo OAuth, confira Autorizar o acesso autônomo aos recursos do Azure Databricks com uma entidade de serviço usando o OAuth. Anote o valor do UUID ou da ID do aplicativo da entidade de serviço e o valor do Segredo do segredo OAuth da entidade de serviço.

2. Dê a essa entidade de serviço acesso ao cluster ou ao warehouse.

   Para conceder acesso à entidade de serviço no seu cluster ou warehouse, confira Permissões de computação ou Gerenciar um SQL warehouse.

Para autenticar o Databricks SQL Connector for Python, use o seguinte trecho de código. Esse snippet de código pressupõe que você tenha definido as seguintes variáveis de ambiente:

• DATABRICKS_SERVER_HOSTNAME definido como o valor Nome de Host do Servidor do seu cluster ou Armazém SQL.
• DATABRICKS_HTTP_PATH, definido para o valor do Caminho HTTP para seu cluster ou SQL Warehouse.
• DATABRICKS_CLIENT_ID, definido como o valor do UUID ou do ID da Aplicação do principal de serviço.
• DATABRICKS_CLIENT_SECRET, definido como o valor do Segredo para o segredo OAuth da entidade de serviço.

Para definir variáveis de ambiente, confira a documentação do sistema operacional.

from databricks.sdk.core import Config, oauth_service_principal
from databricks import sql
import os

server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME")

def credential_provider():
  config = Config(
    host          = f"https://{server_hostname}",
    client_id     = os.getenv("DATABRICKS_CLIENT_ID"),
    client_secret = os.getenv("DATABRICKS_CLIENT_SECRET"))
  return oauth_service_principal(config)

with sql.connect(server_hostname      = server_hostname,
                 http_path            = os.getenv("DATABRICKS_HTTP_PATH"),
                 credentials_provider = credential_provider) as connection:
# ...

Autenticação de token do Microsoft Entra ID

Para usar o Conector do SQL do Databricks para Python com a autenticação de token do Microsoft Entra ID, você deve fornecer ao Conector do SQL do Databricks para Python o token do Microsoft Entra ID. Para criar um token de acesso do Microsoft Entra ID, faça o seguinte:

• Para um usuário do Azure Databricks, você pode usar a CLI do Azure. Confira Obter tokens do Microsoft Entra ID para usuários usando a CLI do Azure.
• Para obter uma entidade de serviço do Microsoft Entra ID, confira Obter um token de acesso do Microsoft Entra ID com a CLI do Azure. Para criar uma entidade de serviço gerenciada do Microsoft Entra ID, consulte Gerenciar entidades de serviço.

Os tokens do Microsoft Entra ID têm um tempo de vida padrão de cerca de 1 hora. Para criar um novo token do Microsoft Entra ID, repita esse processo.

Para autenticar o Databricks SQL Connector for Python, use o seguinte trecho de código. Esse snippet de código pressupõe que você tenha definido as seguintes variáveis de ambiente:

• Defina DATABRICKS_SERVER_HOSTNAME como o valor Nome do host de Servidor do seu cluster ou SQL Warehouse.
• Defina DATABRICKS_HTTP_PATH para Caminho HTTP o valor do seu cluster ou SQL Warehouse.
• Defina DATABRICKS_TOKEN como o token do Microsoft Entra ID.

Para definir variáveis de ambiente, confira a documentação do sistema operacional.

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:
# ...

Autenticação U2M (usuário para computador) do OAuth

As versões 2.7.0 e posteriores do Conector SQL do Databricks para Python suportam a autenticação usuário para computador (U2M) do OAuth. Você também precisa instalar o SDK do Databricks para Python 0.19.0 ou superior, por exemplo, executando pip install databricks-sdk ou python -m pip install databricks-sdk.

Para autenticar o Conector SQL do Databricks para Python com a autenticação OAuth U2M, use o seguinte trecho de código. A autenticação OAuth U2M usa a entrada e consentimento humano em tempo real para autenticar a conta de usuário do Azure Databricks de destino. Esse snippet de código pressupõe que você tenha definido as seguintes variáveis de ambiente:

• Defina DATABRICKS_SERVER_HOSTNAME como o valor Nome do host de Servidor do seu cluster ou SQL Warehouse.
• Defina DATABRICKS_HTTP_PATH para Caminho HTTP o valor do seu cluster ou SQL Warehouse.

Para definir variáveis de ambiente, confira a documentação do sistema operacional.

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 auth_type       = "databricks-oauth") as connection:
# ...

Exemplos

Os exemplos de código a seguir demonstram como usar o Conector de SQL do Databricks para Python para consultar e inserir dados, consultar metadados, gerenciar cursores e conexões e configurar o log.

Esses exemplos de código recuperam os valores de variável de conexão server_hostname, http_path e access_token destas variáveis de ambiente:

• DATABRICKS_SERVER_HOSTNAME, que representa o valor do Nome do host do servidor nos requisitos.
• DATABRICKS_HTTP_PATH, que representa o valor do Caminho HTTP obtido nos requisitos.
• DATABRICKS_TOKEN, que representa seu token de acesso dos requisitos.

Você pode usar outras abordagens para recuperar os valores dessa variável de conexão. Usar variáveis de ambiente é apenas uma abordagem entre muitas.

• Definir agente do usuário
• Consultar dados
• Inserir dados
• Consultar metadados
• Gerenciar cursores e conexões
• Gerenciar arquivos em volumes do Catálogo do Unity
• Configurar o registro em log

Definir agente do usuário

O exemplo de código a seguir demonstra como definir o aplicativo User-Agent product_name para acompanhamento de uso.

from databricks import sql
import os

with sql.connect(server_hostname   = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path         = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token      = os.getenv("DATABRICKS_TOKEN"),
                 _user_agent_entry = "product_name") as connection:
  with connection.cursor() as cursor:
    cursor.execute("SELECT 1 + 1")
    result = cursor.fetchall()

    for row in result:
      print(row)

Consultar dados

O exemplo de código a seguir demonstra como chamar o Conector SQL do Databricks para Python para executar um comando SQL básico em um cluster ou armazém SQL. Esse comando retorna as duas primeiras linhas da tabela trips no esquema samples do catálogo nyctaxi.

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:

  with connection.cursor() as cursor:
    cursor.execute("SELECT * FROM samples.nyctaxi.trips LIMIT 2")
    result = cursor.fetchall()

    for row in result:
      print(row)

Inserir dados

O exemplo a seguir demonstra como inserir pequenas quantidades de dados (milhares de linhas):

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:

  with connection.cursor() as cursor:
    cursor.execute("CREATE TABLE IF NOT EXISTS squares (x int, x_squared int)")

    squares = [(i, i * i) for i in range(100)]
    values = ",".join([f"({x}, {y})" for (x, y) in squares])

    cursor.execute(f"INSERT INTO squares VALUES {values}")

    cursor.execute("SELECT * FROM squares LIMIT 10")

    result = cursor.fetchall()

    for row in result:
      print(row)

Para grandes quantidades de dados, você deve primeiro carregar os dados no armazenamento em nuvem e, em seguida, executar o comando COPY INTO.

Consultar metadados

Há métodos dedicados para recuperar metadados. O exemplo a seguir recupera metadados sobre colunas em uma tabela de exemplo:

from databricks import sql
import os

with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token    = os.getenv("DATABRICKS_TOKEN")) as connection:

  with connection.cursor() as cursor:
    cursor.columns(schema_name="default", table_name="squares")
    print(cursor.fetchall())

Gerenciar cursores e conexões

É uma melhor prática fechar todas as conexões e cursores que não estão mais em uso. Isso libera recursos em clusters do Azure Databricks e em SQL warehouses do Databricks.

Você pode usar um gerenciador de contexto ( a sintaxe with usada nos exemplos anteriores) para gerenciar os recursos ou chamar explicitamente close:

from databricks import sql
import os

connection = sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                         http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                         access_token    = os.getenv("DATABRICKS_TOKEN"))

cursor = connection.cursor()

cursor.execute("SELECT * from range(10)")
print(cursor.fetchall())

cursor.close()
connection.close()

Gerenciar arquivos em volumes do Catálogo do Unity

O Conector do SQL do Databricks permite que você grave arquivos locais em volumes do Catálogo do Unity, baixe arquivos de volumes e exclua arquivos de volumes, conforme mostrado no exemplo a seguir:

from databricks import sql
import os

# For writing local files to volumes and downloading files from volumes,
# you must set the staging_allows_local_path argument to the path to the
# local folder that contains the files to be written or downloaded.
# For deleting files in volumes, you must also specify the
# staging_allows_local_path argument, but its value is ignored,
# so in that case its value can be set for example to an empty string.
with sql.connect(server_hostname            = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                 http_path                  = os.getenv("DATABRICKS_HTTP_PATH"),
                 access_token               = os.getenv("DATABRICKS_TOKEN"),
                 staging_allowed_local_path = "/tmp/") as connection:

  with connection.cursor() as cursor:

    # Write a local file to the specified path in a volume.
    # Specify OVERWRITE to overwrite any existing file in that path.
    cursor.execute(
      "PUT '/temp/my-data.csv' INTO '/Volumes/main/default/my-volume/my-data.csv' OVERWRITE"
    )

    # Download a file from the specified path in a volume.
    cursor.execute(
      "GET '/Volumes/main/default/my-volume/my-data.csv' TO '/tmp/my-downloaded-data.csv'"
    )

    # Delete a file from the specified path in a volume.
    cursor.execute(
      "REMOVE '/Volumes/main/default/my-volume/my-data.csv'"
    )

Configurar o registro em log

O Conector SQL do Databricks usa o módulo padrão de log do Python. Você pode configurar o nível de log como o seguinte:

from databricks import sql
import os, logging

logging.getLogger("databricks.sql").setLevel(logging.DEBUG)
logging.basicConfig(filename = "results.log",
                    level    = logging.DEBUG)

connection = sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
                         http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
                         access_token    = os.getenv("DATABRICKS_TOKEN"))

cursor = connection.cursor()

cursor.execute("SELECT * from range(10)")

result = cursor.fetchall()

for row in result:
   logging.debug(row)

cursor.close()
connection.close()

Testando

Para testar seu código, use estruturas de teste do Python, como pytest. Para testar seu código em condições simuladas sem chamar endpoints da API REST do Azure Databricks ou alterar o estado das suas contas ou espaços de trabalho do Azure Databricks, você pode usar bibliotecas de simulação do Python, como unittest.mock.

Por exemplo, dado o seguinte arquivo chamado helpers.py contendo uma função get_connection_personal_access_token que usa um token de acesso pessoal do Azure Databricks para retornar uma conexão a um workspace do Azure Databricks e uma função select_nyctaxi_trips que usa a conexão para obter o número especificado de linhas de dados da tabela trips no esquema samples do catálogo nyctaxi:

# helpers.py

from databricks import sql
from databricks.sql.client import Connection, List, Row, Cursor

def get_connection_personal_access_token(
  server_hostname: str,
  http_path: str,
  access_token: str
) -> Connection:
  return sql.connect(
    server_hostname = server_hostname,
    http_path = http_path,
    access_token = access_token
  )

def select_nyctaxi_trips(
  connection: Connection,
  num_rows: int
) -> List[Row]:
  cursor: Cursor = connection.cursor()
  cursor.execute(f"SELECT * FROM samples.nyctaxi.trips LIMIT {num_rows}")
  result: List[Row] = cursor.fetchall()
  return result

E dado o seguinte arquivo chamado main.py que chama as funções get_connection_personal_access_token e select_nyctaxi_trips:

# main.py

from databricks.sql.client import Connection, List, Row
import os
from helpers import get_connection_personal_access_token, select_nyctaxi_trips

connection: Connection = get_connection_personal_access_token(
  server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
  http_path = os.getenv("DATABRICKS_HTTP_PATH"),
  access_token = os.getenv("DATABRICKS_TOKEN")
)

rows: List[Row] = select_nyctaxi_trips(
  connection = connection,
  num_rows = 2
)

for row in rows:
  print(row)

O arquivo nomeado test_helpers.py a seguir testa se a função select_nyctaxi_trips retorna a resposta esperada. Em vez de criar uma conexão real com o workspace de destino, esse teste simula um objeto Connection. O teste também simula alguns dados que estão em conformidade com o esquema e os valores que estão nos dados reais. O teste retorna os dados simulados por meio da conexão simulada e verifica se um dos valores das linhas de dados simuladas corresponde ao valor esperado.

# test_helpers.py

import pytest
from databricks.sql.client import Connection, List, Row
from datetime import datetime
from helpers import select_nyctaxi_trips
from unittest.mock import create_autospec

@pytest.fixture
def mock_data() -> List[Row]:
  return [
    Row(
      tpep_pickup_datetime = datetime(2016, 2, 14, 16, 52, 13),
      tpep_dropoff_datetime = datetime(2016, 2, 14, 17, 16, 4),
      trip_distance = 4.94,
      fare_amount = 19.0,
      pickup_zip = 10282,
      dropoff_zip = 10171
    ),
    Row(
      tpep_pickup_datetime = datetime(2016, 2, 4, 18, 44, 19),
      tpep_dropoff_datetime = datetime(2016, 2, 4, 18, 46),
      trip_distance = 0.28,
      fare_amount = 3.5,
      pickup_zip = 10110,
      dropoff_zip = 10110
    )
  ]

def test_select_nyctaxi_trips(mock_data: List[Row]):
  # Create a mock Connection.
  mock_connection = create_autospec(Connection)

  # Set the mock Connection's cursor().fetchall() to the mock data.
  mock_connection.cursor().fetchall.return_value = mock_data

  # Call the real function with the mock Connection.
  response: List[Row] = select_nyctaxi_trips(
    connection = mock_connection,
    num_rows = 2)

  # Check the value of one of the mocked data row's columns.
  assert response[1].fare_amount == 3.5

Como a função select_nyctaxi_trips contém uma instrução SELECT e, portanto, não altera o estado da tabela trips, a simulação não é absolutamente necessária neste exemplo. No entanto, a simulação permite que você execute rapidamente seus testes sem esperar que uma conexão real seja feita com o workspace. Além disso, o mocking permite executar testes simulados várias vezes para funções que podem alterar o estado de uma tabela, como INSERT INTO, UPDATE e DELETE FROM.

Referência da API

• Pacote
• Módulo
• Classes
• Connection classe
• Cursor classe
• Row classe
• Conversões de tipo

Pacote

databricks-sql-connector

Uso: pip install databricks-sql-connector

Consulte também databricks-sql-connector no PyPI (Índice de Pacotes Python).

Módulo

databricks.sql

Uso: from databricks import sql

Classes

As classes selecionadas incluem o seguinte:

Classe Connection

Para criar um objeto Connection, chame o método databricks.sql.connect com os seguintes parâmetros:

Os métodos Connection selecionados incluem o seguinte:

Classe Cursor

Para criar um objeto Cursor, chame o método Connection da classe cursor.

Os atributos Cursor selecionados incluem o seguinte:

Os métodos Cursor selecionados incluem o seguinte:

Classe Row

A classe de linha é uma estrutura de dados do tipo tupla que representa uma linha de resultado individual. Se a linha contiver uma coluna com o nome "my_column", você poderá acessar o campo "my_column" de row via row.my_column. Você também pode usar indicies numéricos para acessar campos, por exemplo row[0]. Se o nome da coluna não for permitido como um nome de método de atributo (por exemplo, ele começa com um dígito), você poderá acessar o campo como row["1_my_column"].

Desde a versão 1.0

Os métodos Row selecionados incluem:

| asDict

Retorna uma representação de dicionário da linha, que é indexada por nomes de campo. Se houver nomes de campo duplicados, um dos campos duplicados (mas apenas um) será retornado no dicionário. Qual campo duplicado é retornado não está definido.

Sem parâmetros.

Retorna um dict de campos. |

Conversões de tipo

A tabela a seguir mapeia os tipos de dados de SQL do Apache Spark para seus equivalentes de tipo de dados Python.

Solução de problemas

mensagem tokenAuthWrapperInvalidAccessToken: Invalid access token

Problema: ao executar seu código, você vê uma mensagem semelhante a Error during request to server: tokenAuthWrapperInvalidAccessToken: Invalid access token.

Causa possível: o valor passado para access_token não é um token de acesso pessoal válido do Azure Databricks.

Correção recomendada: verifique se o valor passado para access_token está correto e tente novamente.

mensagem gaierror(8, 'nodename nor servname provided, or not known')

Problema: ao executar seu código, você vê uma mensagem semelhante a Error during request to server: gaierror(8, 'nodename nor servname provided, or not known').

Causa possível: o valor passado para server_hostname não é o nome do host correto.

Correção recomendada: verifique se o valor passado para server_hostname está correto e tente novamente.

Para obter mais informações sobre como localizar o nome do host do servidor, confira Obter detalhes de conexão para um recurso de computação do Azure Databricks.

mensagem IpAclError

Problema: quando você executar seu código, verá a mensagem Error during request to server: IpAclValidation ao tentar usar o conector em um notebook do Azure Databricks.

Causa possível: você pode ter a listagem de permitir IP habilitada para o workspace do Azure Databricks. Com a listagem de permissão de IP, as conexões de clusters Spark de volta para o plano de controle não são permitidas por padrão.

Correção recomendada: peça ao administrador para adicionar a sub-rede do plano de computação à lista de permissões de IP.

Recursos adicionais

Para obter mais informações, consulte:

• O repositório do Conector do Databricks SQL para Python no GitHub
• Tipos de dados
• Tipos integrados (para bool, bytearray, float, int, e str) no site do Python
• datetime (para datetime.date e datatime.datetime) no site do Python
• decimal (para decimal.Decimal) no site do Python
• Constantes integrados (para NoneType) no site do Python