|
| 1 | +--- |
| 2 | +sidebar_position: 1 |
| 3 | +--- |
| 4 | + |
| 5 | +# Read a Lakebase table |
| 6 | + |
| 7 | +This app connects to a [Databricks Lakebase](https://docs.databricks.com/aws/en/oltp/) OLTP database instance and reads the first 100 rows from any table. Provide the instance name, database, schema, and table name. |
| 8 | + |
| 9 | +## Code snippet |
| 10 | + |
| 11 | +```python title="app.py" |
| 12 | +import os |
| 13 | +import pandas as pd |
| 14 | +import psycopg |
| 15 | +from databricks.sdk import WorkspaceClient |
| 16 | +import streamlit as st |
| 17 | + |
| 18 | + |
| 19 | +w = WorkspaceClient() |
| 20 | + |
| 21 | + |
| 22 | +def get_connection(host: str, database: str, user: str) -> psycopg.Connection: |
| 23 | + """Get a connection to the Lakebase database using OAuth token.""" |
| 24 | + token = w.config.oauth_token().access_token |
| 25 | + |
| 26 | + return psycopg.connect( |
| 27 | + host=host, |
| 28 | + port=5432, |
| 29 | + dbname=database, |
| 30 | + user=user, |
| 31 | + password=token, |
| 32 | + sslmode="require", |
| 33 | + ) |
| 34 | + |
| 35 | + |
| 36 | +def query_df(host: str, database: str, user: str, sql: str) -> pd.DataFrame: |
| 37 | + """Execute a SQL query and return results as a DataFrame.""" |
| 38 | + conn = get_connection(host, database, user) |
| 39 | + try: |
| 40 | + with conn.cursor() as cur: |
| 41 | + cur.execute(sql) |
| 42 | + if not cur.description: |
| 43 | + return pd.DataFrame() |
| 44 | + |
| 45 | + cols = [d.name for d in cur.description] |
| 46 | + rows = cur.fetchall() |
| 47 | + return pd.DataFrame(rows, columns=cols) |
| 48 | + finally: |
| 49 | + conn.close() |
| 50 | + |
| 51 | + |
| 52 | +# Get connection parameters from environment variables (set by Databricks Apps) |
| 53 | +# or fall back to manual configuration |
| 54 | +host = os.getenv("PGHOST") |
| 55 | +database = os.getenv("PGDATABASE") |
| 56 | +user = os.getenv("PGUSER") |
| 57 | + |
| 58 | +if not all([host, database, user]): |
| 59 | + # Manual configuration if environment variables are not set |
| 60 | + instance_name = "your_instance_name" |
| 61 | + database = "databricks_postgres" |
| 62 | + user = w.config.client_id or w.current_user.me().user_name |
| 63 | + host = w.database.get_database_instance(name=instance_name).read_write_dns |
| 64 | + |
| 65 | +# Query table |
| 66 | +schema = "public" |
| 67 | +table = "your_table_name" |
| 68 | +df = query_df(host, database, user, f"SELECT * FROM {schema}.{table} LIMIT 100") |
| 69 | +st.dataframe(df) |
| 70 | +``` |
| 71 | + |
| 72 | +:::tip |
| 73 | + |
| 74 | +Add your Lakebase instance as an App resource to automatically configure connection parameters via environment variables. See the [Lakebase resource documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/lakebase) for details. |
| 75 | + |
| 76 | +::: |
| 77 | + |
| 78 | +:::warning |
| 79 | + |
| 80 | +Tokens expire periodically; this app refreshes on each new connection and enforces TLS (sslmode=require). |
| 81 | + |
| 82 | +::: |
| 83 | + |
| 84 | +## Resources |
| 85 | + |
| 86 | +- [Lakebase](https://docs.databricks.com/aws/en/oltp/) database instance (Postgres). |
| 87 | +- An existing Postgres database, schema, and table with data. |
| 88 | + |
| 89 | +## Permissions |
| 90 | + |
| 91 | +Add the Lakebase instance as an [**App resource**](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/lakebase) to automatically configure permissions and environment variables (`PGHOST`, `PGDATABASE`, `PGUSER`, etc.). |
| 92 | + |
| 93 | +Alternatively, manually create a Postgres role for the service principal. See [this guide](https://docs.databricks.com/aws/en/oltp/pg-roles?language=PostgreSQL#create-postgres-roles-and-grant-privileges-for-databricks-identities). |
| 94 | + |
| 95 | +Example grants for read access: |
| 96 | + |
| 97 | +```sql |
| 98 | +GRANT CONNECT ON DATABASE databricks_postgres TO "<service-principal-id>"; |
| 99 | +GRANT USAGE ON SCHEMA public TO "<service-principal-id>"; |
| 100 | +GRANT SELECT ON TABLE your_table_name TO "<service-principal-id>"; |
| 101 | +``` |
| 102 | + |
| 103 | +[This guide](https://docs.databricks.com/aws/en/oltp/query/sql-editor#create-a-new-query) shows you how to query your Lakebase. |
| 104 | + |
| 105 | +## Dependencies |
| 106 | + |
| 107 | +- [Databricks SDK](https://pypi.org/project/databricks-sdk/) - `databricks-sdk>=0.60.0` |
| 108 | +- [Psycopg](https://pypi.org/project/psycopg/) - `psycopg[binary]` |
| 109 | +- [Pandas](https://pypi.org/project/pandas/) - `pandas` |
| 110 | +- [Streamlit](https://pypi.org/project/streamlit/) - `streamlit` |
| 111 | + |
| 112 | +```python title="requirements.txt" |
| 113 | +databricks-sdk>=0.60.0 |
| 114 | +pandas |
| 115 | +streamlit |
| 116 | +psycopg[binary] |
| 117 | +``` |
| 118 | + |
0 commit comments