Skip to navigation

PostgreSQL Connector

View as Markdown

There are two ways to use PostgreSQL data: sample it in place with PostgresSource, or read a table into memory with read().

PostgresSource requires the relational engine (the [relational] extra), because it was carried over from before the client gained a standalone Postgres connector — the engine provides the sampler.

Installation

pip install "structured-data-models-client[relational,postgres]"

Sample in place with PostgresSource

import psycopg
from sdm_client import PredictiveQuery, Schema
from sdm_connectors import PostgresSource
connection = psycopg.connect(
host="db.example.com", dbname="analytics", user="reader", password="...",
)
source = PostgresSource(connection, schema="public")
schema = Schema.from_source(source, {"users": "users", "orders": "orders"})
request = PredictiveQuery(
"PREDICT COUNT(orders.*, 0, 30, days) FOR EACH users.user_id", schema
).materialize(indices=[7, 12, 3])

tables maps the names your query uses to the names in PostgreSQL; leave it out to describe every table in the schema. The connection is yours — close it when you are done sampling. Because PostgreSQL’s own primary-key constraint is enforced by the database, Schema.from_source trusts a declared key there and skips the uniqueness scan it runs against a warehouse that does not enforce one.

Read into memory instead

from sdm_connectors import read
from sdm_client import Schema
users = read("postgres", host="db.example.com", dbname="analytics", user="reader", password="...", table="users")
orders = read("postgres", host="db.example.com", dbname="analytics", user="reader", password="...", table="orders")
schema = Schema.from_data({"users": users, "orders": orders})

Standard PG* environment variables (PGHOST, PGPORT, PGDATABASE, PGUSER, PGPASSWORD, PGSSLMODE) fill in any of host, port, dbname, user, password, sslmode you omit. A libpq conninfo string or URI also works: read("postgres", conninfo="postgresql://reader@db.example.com/analytics", table="orders").

For a Databricks Lakebase database specifically, use the Lakebase connector instead — it authenticates with a short-lived token rather than a static password.

Review the inferred schema and relationships before predicting — see Build a Schema.