Pular para o conteúdo principal

Usar um Notebook para acessar uma instância de banco de dados

info

Visualização

Esse recurso está em Public Preview nas seguintes regiões: us-east-1, us-west-2, eu-west-1, ap-southeast-1, ap-southeast-2, eu-central-1, us-east-2, ap-south-1.

Esta página contém exemplos de código que mostram como acessar a instância do banco de dados do Lakebase por meio do Databricks Notebook e executar consultas usando Python e Scala.

Os exemplos abrangem diferentes estratégias de conexão para se adequar a diferentes casos de uso:

  • Conexão única : usada para scripts simples em que uma única conexão de banco de dados é aberta, usada e fechada.
  • Conexão pool : Usada para cargas de trabalho de alta simultaneidade, em que é mantida uma pool de conexões reutilizáveis.
  • Tokens rotativos M2M OAuth: Usa OAuth tokens de curta duração, atualizados automaticamente para autenticação.

Os exemplos a seguir geram credenciais seguras de forma programática. Evite colocar credenciais diretamente em um Notebook. A Databricks recomenda o uso de um dos seguintes métodos seguros:

  • Armazene as senhas do Postgres nos segredos do Databricks.
  • Gerar tokens OAuth usando o M2M OAuth.

Antes de começar​

Certifique-se de atender aos seguintes requisitos antes de acessar sua instância de banco de dados:

  • O senhor tem uma função Postgres correspondente para log in na instância do banco de dados. Consulte gerenciar funções do Postgres.
  • Sua função Postgres recebe as permissões necessárias para acessar o banco de dados, esquema ou tabela.
  • Você pode se autenticar na instância do banco de dados. Se o senhor precisar obter manualmente um token OAuth para a instância do banco de dados, consulte Autenticação em uma instância do banco de dados.

Python​

O Databricks Python SDK pode ser usado para obter tokens OAuth para uma respectiva instância de banco de dados.

Conecte-se à instância do banco de dados a partir de um Notebook Databricks usando a seguinte Python biblioteca:

  • psycopg2
  • psycopg3
  • SQLAlchemy

psycopg2​

Os exemplos de código demonstram uma única conexão e o uso de um pool de conexões. Para saber mais sobre como obter a instância do banco de dados e as credenciais de forma programática, consulte como obter tokens OAuth usando o Python SDK .

Python
import psycopg2

from databricks.sdk import WorkspaceClient
import uuid

w = WorkspaceClient()

instance_name = "<YOUR INSTANCE>"

instance = w.database.get_database_instance(name=instance_name)
cred = w.database.generate_database_credential(request_id=str(uuid.uuid4()), instance_names=[instance_name])

# Connection parameters
conn = psycopg2.connect(
    host = instance.read_write_dns,
    dbname = "databricks_postgres",
    user = "<YOUR USER>",
    password = cred.token,
    sslmode = "require"
)

# Execute query
with conn.cursor() as cur:
    cur.execute("SELECT version()")
    version = cur.fetchone()[0]
    print(version)
conn.close()

psycopg3​

O exemplo de código demonstra o uso de um pool de conexões com um M2M OAuth rotativo. Ele usa generate_database_credential(). Para saber mais sobre como obter a instância do banco de dados e as credenciais de forma programática, consulte como obter tokens OAuth usando o Python SDK .

Python
%pip install "psycopg[binary,pool]"
Python
from databricks.sdk import WorkspaceClient
import uuid

import psycopg
import string
from psycopg_pool import ConnectionPool

w = WorkspaceClient()

class CustomConnection(psycopg.Connection):
    global w
    def __init__(self, *args, **kwargs):
        # Call the parent class constructor
        super().__init__(*args, **kwargs)

    @classmethod
    def connect(cls, conninfo='', **kwargs):
        # Append the new password to kwargs
        cred = w.database.generate_database_credential(request_id=str(uuid.uuid4()), instance_names=[instance_name])
        kwargs['password'] = cred.token

        # Call the superclass's connect method with updated kwargs
        return super().connect(conninfo, **kwargs)


username = "<YOUR USER>"
instance_name = "<YOUR INSTANCE>"
instance = w.database.get_database_instance(name=instance_name)
host = instance.read_write_dns
port = 5432
database = "databricks_postgres"

pool = ConnectionPool(
    conninfo=f"dbname={database} user={username} host={host}",
    connection_class=CustomConnection,
    min_size=1,
    max_size=10,
    open=True
)

with pool.connection() as conn:
    with conn.cursor() as cursor:
        cursor.execute("SELECT version()")
        for record in cursor:
            print(record)

SQLAlchemy​

Os exemplos de código demonstram uma única conexão e o uso de uma conexão pool com tokens rotativos M2M OAuth. Para saber mais sobre como obter a instância do banco de dados e as credenciais de forma programática, consulte como obter tokens OAuth usando o Python SDK .

Python
%pip install sqlalchemy==1.4 psycopg[binary]
Python
from sqlalchemy import create_engine, text

from databricks.sdk import WorkspaceClient
import uuid

w = WorkspaceClient()

instance_name = "<YOUR INSTANCE>"

instance = w.database.get_database_instance(name=instance_name)
cred = w.database.generate_database_credential(request_id=str(uuid.uuid4()), instance_names=[instance_name])

user = "<YOUR USER>"
host = instance.read_write_dns
port = 5432
database = "databricks_postgres"
password = cred.token

connection_pool = create_engine(f"postgresql://{user}:{password}@{host}:{port}/{database}?sslmode=require")

with connection_pool.connect() as conn:
    result = conn.execute(text("SELECT version()"))
    for row in result:
        print(f"Connected to PostgreSQL database. Version: {row}")

Scala​

Os exemplos de código mostram como obter programaticamente a instância e as credenciais do banco de dados e como se conectar a uma instância do banco de dados usando uma única conexão ou um pool de conexões.

Etapa 1: Use o endereço Databricks Java SDK para obter tokens OAuth​

Para obter detalhes sobre como obter a instância do banco de dados e as credenciais de forma programática, consulte como obter tokens OAuth usando o Java SDK .

Etapa 2: conectar-se a uma instância de banco de dados​

Scala
import java.sql.{Connection, DriverManager, ResultSet, Statement}

Class.forName("org.postgresql.Driver")

val user = "<YOUR USER>"
val host = instance.getName()
val port = "5432"
val database = "databricks_postgres"
val password = cred.getToken()

val url = f"jdbc:postgresql://${host}:${port}/${database}"

val connection = DriverManager.getConnection(url, user, password)
println("Connected to PostgreSQL database!")

val statement = connection.createStatement()
val resultSet = statement.executeQuery("SELECT version()")

if (resultSet.next()) {
    val version = resultSet.getString(1)
    println(s"PostgreSQL version: $version")
}
Última atualização em