Skip to main content

LakeBase with Python Client

Connecting from a Python client (with Onehouse Catalog)​

Use this when your LakeBase Cluster is configured with Onehouse Catalog. Authentication is handled via the username/password shown on the cluster's connection details page, and authorization is controlled by the Onehouse's EKS node role.

Step 1: Install psycopg2​

pip install psycopg2-binary

Step 2: Get Connection Details​

From the Onehouse console, open the Clusters page, click into your LakeBase cluster, and copy the Endpoint URL, username, and password. Refer to Get Connection Details for more.

If VPN is not enabled

If VPN is not enabled, configure a bastion host to reach the cluster. Refer to Connect to your VPC.

Step 3: Connect and Run a Query​

import psycopg2

conn = psycopg2.connect(
    host="<cluster-host>",
    port=5432,
    dbname="postgres",
    user="<username>",
    password="<password>",
    connect_timeout=300,
)
conn.autocommit = True

with conn.cursor() as cur:
    cur.execute("SELECT 1")
    print(cur.fetchall())

conn.close()

Connecting from a Python client (with Glue Catalog)​

Use this when your LakeBase Cluster is configured with Glue Catalog (without Lake Formation). The connection pattern is identical to the Onehouse Catalog case above — authentication uses the username/password from the cluster page, and authorization is controlled by the Onehouse's EKS node role.

Follow the same Onehouse Catalog steps: use psycopg2 with the credentials from the cluster connection details page.

Connecting from a Python client (with Glue Catalog with Lake Formation)​

Use this when your LakeBase Cluster is configured with Glue Catalog and Lake Formation. Authentication is delegated to your corporate identity provider (IdP) via the Onehouse LakeBase JDBC driver, and authorization is enforced by Lake Formation.

Since the browser-based authentication flow lives in the JDBC driver, the Python client connects through the JVM using JayDeBeApi + JPype.

Step 1: Install Dependencies​

pip install JayDeBeApi JPype1

You also need a Java runtime (JRE 8 or later) available on the machine running the script.

Step 2: Download the Onehouse LakeBase JDBC Driver​

Download the LakeBase JDBC driver and note the local path to the JAR. The driver class is ai.onehouse.jdbc.OnehouseDriver.

Step 3: Get Connection Details​

From the Onehouse console, open the Clusters page, click into your LakeBase cluster, and copy the Endpoint URL. Refer to Get Connection Details for more.

If VPN is not enabled

If VPN is not enabled, configure a bastion host to reach the cluster. Refer to Connect to your VPC.

Step 4: Connect and Run a Query​

Set the IdP-specific connection properties for your provider. Refer to Supported Identity Providers for the full list of properties.

The example below uses Azure AD OAuth2:

import jaydebeapi
import jpype
import jpype.imports

JAR_PATH = "/path/to/lakebase-driver.jar"
DRIVER_CLASS = "ai.onehouse.jdbc.OnehouseDriver"

HOST = "<cluster-host>"
PORT = "5432"
DBNAME = "postgres"

# Start the JVM with the driver on the classpath
if not jpype.isJVMStarted():
    jpype.startJVM(classpath=[JAR_PATH])

from java.util import Properties
from java.sql import DriverManager

jpype.JClass(DRIVER_CLASS)  # load driver

props = Properties()
# Placeholder user/password — the driver authenticates via the IdP
props.setProperty("user", "user")
props.setProperty("password", "user")

# Browser-based IdP authentication
props.setProperty("browserAuth", "true")
props.setProperty("azureOAuthTenantId", "<tenant-id>")
props.setProperty("azureOAuthClientId", "<client-id>")
props.setProperty("azureOAuthClientSecret", "<client-secret>")
props.setProperty("azureOAuthIamRole",
    "arn:aws:iam::<account>:role/<onehouse-lakebase-role>")

print("Browser will open for OAuth authentication...")

# Initialize JayDeBeApi type converters
if jaydebeapi._jdbc_name_to_const is None:
    types = jpype.java.sql.Types
    types_map = {}
    for field in types.class_.getFields():
        if jpype.java.lang.reflect.Modifier.isStatic(field.getModifiers()):
            types_map[field.getName()] = field.get(None)
    jaydebeapi._init_types(types_map)

jdbc_url = f"jdbc:postgresql://{HOST}:{PORT}/{DBNAME}"
java_conn = DriverManager.getConnection(jdbc_url, props)

# Wrap in a DB-API cursor
conn = jaydebeapi.Connection(java_conn, jaydebeapi._converters)

with conn.cursor() as cur:
    cur.execute("SELECT 1")
    print(cur.fetchall())

conn.close()

On DriverManager.getConnection(), the driver opens a browser window for the IdP login, starts a local callback server (default port 8888) to receive the response, and then completes the PostgreSQL connection.

Step 5: (Optional) Specify an IAM Role​

When using Lake Formation for access control, you can optionally specify an IAM role to assume for querying. If you don't specify a role, the first IAM role from your authentication token is used by default.

To specify a role explicitly, add the options property to the Java Properties object:

props.setProperty(
    "options",
    "-c onehouse.lakebase_iam_role=arn:aws:iam::<account>:role/<role-name>",
)

Or include it in the JDBC URL:

jdbc_url = (
    f"jdbc:postgresql://{HOST}:{PORT}/{DBNAME}"
    "?options=-c onehouse.lakebase_iam_role=arn:aws:iam::<account>:role/<role-name>"
)