使用笔记本访问数据库实例

重要

此功能在以下区域中为公共预览版：westus、westus2、eastus、eastus2、centralus、southcentralus、northeurope、westeurope、australiaeast、brazilsouth、canadacentral、centralindia、southeastasia、uksouth。

本页包含代码示例，演示如何通过 Azure Databricks 笔记本访问 Lakebase 数据库实例，以及如何使用 Python 和 Scala 运行查询。

这些示例涵盖不同的连接策略，以适应不同的用例：

单一连接：用于打开、使用和关闭单一数据库连接的简单脚本。
连接池：用于处理高并发工作负载，在其中维护了一个可重用连接的集合。
轮换 M2M OAuth 令牌：使用生存期较短、自动刷新的 OAuth 令牌进行身份验证。

以下示例以编程方式生成安全凭据。避免直接将凭据放入笔记本中。 Databricks 建议使用以下安全方法之一：

将 Postgres 密码存储在 Azure Databricks 机密中。
使用 M2M OAuth 生成 OAuth 令牌。

在您开始之前

在访问数据库实例之前，请确保满足以下要求：

你拥有一个相应的 Postgres 角色，专用于登录数据库实例。请参阅 “管理 Postgres 角色”。
你的 Postgres 角色已经被授予了访问数据库、架构或表所需的权限。
可以对数据库实例进行身份验证。如果必须为数据库实例手动获取 OAuth 令牌，请参阅对数据库实例进行身份验证。

警告

如果使用专用链接，则需要使用单个用户群集。

Python

Azure Databricks Python SDK 可用于获取相应数据库实例的 OAuth 令牌。

使用以下 Python 库从 Azure Databricks 笔记本连接到数据库实例：

psycopg2
psycopg3
SQLAlchemy

先决条件

运行以下代码示例之前，请将用于 Python 的 Databricks SDK 升级到版本 0.61.0 或更高版本，然后重启 Python。

%pip install databricks-sdk>=0.61.0
%restart_python

`psycopg2`

代码示例演示单个连接和使用连接池。有关如何以编程方式获取数据库实例和凭据的详细信息，请参阅如何使用 Python SDK 获取 OAuth 令牌。

单个连接

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()

连接池

import psycopg2
from psycopg2 import sql, pool
from pyspark.sql.functions import udf

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])

# Create a connection pool
connection_pool = psycopg2.pool.ThreadedConnectionPool(
    minconn = 1,  # Minimum number of connections in the pool
    maxconn = 10,  # Maximum number of connections in the pool
    user = "<YOUR USER>",
    password = cred.token,
    host = instance.read_write_dns,
    port = '5432',
    database = 'databricks_postgres'
)
if connection_pool:
    print("Connection pool created successfully")


def executeWithPgConnection(execFn):
    connection = None
    try:
        # Get a connection from the pool
        connection = connection_pool.getconn()
        if connection:
            print("Successfully received a connection from the pool")
            execFn(connection)
    finally:
        # Release the connection back to the pool
        if connection:
            connection_pool.putconn(connection)
            print("Connection returned to the pool")


def printVersion(connection):
    cursor = connection.cursor()
    cursor.execute("SELECT version()")
    version = cursor.fetchone()
    print(f"Connected to PostgreSQL database. Version: {version}")

executeWithPgConnection(printVersion)

`psycopg3`

该代码示例演示如何将连接池与旋转的 M2M OAuth 配合使用。它使用 generate_database_credential()。有关如何以编程方式获取数据库实例和凭据的详细信息，请参阅如何使用 Python SDK 获取 OAuth 令牌。

%pip install "psycopg[binary,pool]"

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`

代码示例演示单个连接，以及将连接池与旋转的 M2M OAuth 令牌配合使用。有关如何以编程方式获取数据库实例和凭据的详细信息，请参阅如何使用 Python SDK 获取 OAuth 令牌。

单个连接

%pip install sqlalchemy==1.4 psycopg[binary]

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}")

连接池和轮换 M2M OAuth

%pip install sqlalchemy==1.4 psycopg[binary]

from databricks.sdk import WorkspaceClient
import uuid
import time

from sqlalchemy import create_engine, text, event

w = WorkspaceClient()

instance_name = "<YOUR INSTANCE>"

instance = w.database.get_database_instance(name=instance_name)

username = "<YOUR USER>"
host = instance.read_write_dns
port = 5432
database = "databricks_postgres"

# sqlalchemy setup + function to refresh the OAuth token that is used as the Postgres password every 15 minutes.
connection_pool = create_engine(f"postgresql+psycopg2://{username}:@{host}:{port}/{database}")
postgres_password = None
last_password_refresh = time.time()

@event.listens_for(connection_pool, "do_connect")
def provide_token(dialect, conn_rec, cargs, cparams):
    global postgres_password, last_password_refresh, host

    if postgres_password is None or time.time() - last_password_refresh > 900:
        print("Refreshing PostgreSQL OAuth token")
        cred = w.database.generate_database_credential(request_id=str(uuid.uuid4()), instance_names=[instance_name])
        postgres_password = cred.token
        last_password_refresh = time.time()

    cparams["password"] = postgres_password

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

代码示例演示如何以编程方式获取数据库实例和凭据，以及如何使用单个连接或连接池连接到数据库实例。

步骤 1：使用 Azure Databricks Java SDK 获取 OAuth 令牌

有关如何以编程方式获取数据库实例和凭据的详细信息，请参阅如何使用 Java SDK 获取 OAuth 令牌。

步骤 2：连接到数据库实例

单个连接

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")
}

连接池

import com.zaxxer.hikari.{HikariConfig, HikariDataSource}
import java.sql.Connection

// Configure HikariCP
val config = new HikariConfig()
config.setJdbcUrl("jdbc:postgresql://instance.getName():5432/databricks_postgres")
config.setUsername("<YOUR USER>")
config.setPassword(cred.getToken())
config.setMaximumPoolSize(10)

// Create a data source
val dataSource = new HikariDataSource(config)

// Function to get a connection and execute a query
def runQuery(): Unit = {
  var connection: Connection = null
  try {
    // Get a connection from the pool
    connection = dataSource.getConnection()

    // Create a statement
    val statement = connection.createStatement()

    // Execute a query
    val resultSet = statement.executeQuery("SELECT version() AS v;")

    // Process the result set
    while (resultSet.next()) {
      val v = resultSet.getString("v")
      println(s"*******Connected to PostgreSQL database. Version: $v")
    }
  } catch {
    case e: Exception => e.printStackTrace()
  } finally {
    // Close the connection which returns it to the pool
    if (connection != null) connection.close()
  }
}

// Run the query
runQuery()

// Close the data source
dataSource.close()

反馈

此页面是否有帮助？

Last updated on 2025-12-12

通过

使用笔记本访问数据库实例

在您开始之前

Python

先决条件

psycopg2

单个连接

连接池

psycopg3

SQLAlchemy

单个连接

连接池和轮换 M2M OAuth

Scala

步骤 1：使用 Azure Databricks Java SDK 获取 OAuth 令牌

步骤 2：连接到数据库实例

单个连接

连接池

反馈

其他资源

`psycopg2`

`psycopg3`

`SQLAlchemy`