Skip to main content

Read a Lakebase table

This app connects to a Databricks Lakebase OLTP database instance and reads the first 100 rows from any table. Provide the instance name, database, schema, and table name.

Code snippet​

app.py
import os
import pandas as pd
import psycopg
from databricks.sdk import WorkspaceClient
import streamlit as st


w = WorkspaceClient()


def get_connection(host: str, database: str, user: str) -> psycopg.Connection:
    """Get a connection to the Lakebase database using OAuth token."""
    token = w.config.oauth_token().access_token
    
    return psycopg.connect(
        host=host,
        port=5432,
        dbname=database,
        user=user,
        password=token,
        sslmode="require",
    )


def query_df(host: str, database: str, user: str, sql: str) -> pd.DataFrame:
    """Execute a SQL query and return results as a DataFrame."""
    conn = get_connection(host, database, user)
    try:
        with conn.cursor() as cur:
            cur.execute(sql)
            if not cur.description:
                return pd.DataFrame()
            
            cols = [d.name for d in cur.description]
            rows = cur.fetchall()
            return pd.DataFrame(rows, columns=cols)
    finally:
        conn.close()


# Get connection parameters from environment variables (set by Databricks Apps)
# or fall back to manual configuration
host = os.getenv("PGHOST")
database = os.getenv("PGDATABASE")
user = os.getenv("PGUSER")

if not all([host, database, user]):
    # Manual configuration if environment variables are not set
    instance_name = "your_instance_name"
    database = "databricks_postgres"
    user = w.config.client_id or w.current_user.me().user_name
    host = w.database.get_database_instance(name=instance_name).read_write_dns

# Query table
schema = "public"
table = "your_table_name"
df = query_df(host, database, user, f"SELECT * FROM {schema}.{table} LIMIT 100")
st.dataframe(df)
tip

Add your Lakebase instance as an App resource to automatically configure connection parameters via environment variables. See the Lakebase resource documentation for details.

warning

Tokens expire periodically; this app refreshes on each new connection and enforces TLS (sslmode=require).

Resources​

  • Lakebase database instance (Postgres).
  • An existing Postgres database, schema, and table with data.

Permissions​

Add the Lakebase instance as an App resource to automatically configure permissions and environment variables (PGHOST, PGDATABASE, PGUSER, etc.).

Alternatively, manually create a Postgres role for the service principal. See this guide.

Example grants for read access:

GRANT CONNECT ON DATABASE databricks_postgres TO "<service-principal-id>";
GRANT USAGE ON SCHEMA public TO "<service-principal-id>";
GRANT SELECT ON TABLE your_table_name TO "<service-principal-id>";

This guide shows you how to query your Lakebase.

Dependencies​

requirements.txt
databricks-sdk>=0.60.0
pandas
streamlit
psycopg[binary]