Production Lakebase Patterns for Databricks Apps: OAuth Token Refresh and Async Pooling

Summary

• Token refresh: proactively refresh and recreate connection pools before the 60-minute expiry.
• Pool watchdog: background process gracefully replaces connection pools with custom token refresh.
• Concurrent writes: use upsert patterns with idempotency keys for safe agent writes at scale.

Production Lakebase Patterns for Databricks Apps: OAuth Token Refresh and Async Pooling

Talk by Irvin Umaña, Lead Scale Solution Engineer, Databricks
A Databricks App backed by Lakebase, the PostgreSQL database on Databricks, can work perfectly in a notebook and then fail with an authentication error at the 61st minute, because OAuth tokens are valid for one hour and Postgres checks expiration only at login. The fix is to refresh tokens proactively and recreate the connection pool before credentials expire. This video shares four production patterns from an energy sector knowledge assistant: a psycopg3 async pool with a custom connection class that fetches fresh tokens, a background pool watchdog that swaps pools gracefully, LRU caching with TTL, and safe upserts under concurrent agent writes.
This lightning talk shares four production patterns from an energy sector knowledge assistant: a psycopg3 async connection pool with a custom connection class that fetches fresh tokens, a background pool watchdog that swaps in a new pool and shuts the old one down gracefully, LRU caching with TTL for high-frequency thread lookups, and upsert patterns using ON CONFLICT DO UPDATE and idempotency keys for concurrent agent writes. It also covers sizing pools for Databricks Apps concurrency, keeping warm connections when Lakebase scales to zero, and hard limits such as the 24-hour inactivity timeout and three-day maximum connection lifetime.

Chapters

• 00:00 Introduction: Lakebase in production apps
• 01:58 The 61st minute authentication error
• 05:55 Four production patterns overview
• 07:48 OAuth token expiry and the connection pool trap
• 09:01 Custom psycopg3 connection class for fresh tokens
• 10:13 Pool watchdog: recreate connection pools before tokens expire
• 11:59 Async pool sizing and scale to zero
• 13:18 LRU cache with TTL for hot thread lookups
• 14:34 Upserts and race conditions with concurrent agent writes
• 17:05 Real-world energy knowledge assistant architecture
• 18:58 Key takeaways

FAQs

Why do Lakebase apps fail after 61 minutes?

OAuth tokens are valid for one hour, and Postgres only checks token expiration at login. The fix is to refresh tokens proactively and recreate the connection pool before the 60-minute mark to prevent authentication errors.

How do I handle concurrent writes to Lakebase?

Use upsert patterns with ON CONFLICT DO UPDATE and idempotency keys to safely handle concurrent agent writes without race conditions or duplicate data.

What's a pool watchdog?

A background process that periodically swaps in a new connection pool before tokens expire, then gracefully shuts down the old pool. This prevents the sudden 61st-minute authentication failure.

Can Lakebase apps handle many concurrent users?

Yes, if you size the psycopg3 async pool for your Databricks Apps concurrency level. The video covers real-world patterns from an energy sector knowledge assistant handling concurrent agents.

Full transcript

00:00 Introduction: Lakebase in production apps

Good morning everyone. Welcome to the lightning talks at the Data and AI Summit 2026. My name is Irwin Umanya, I'm a lead engineer for scalable solutions at Databricks. I work on a shared technical services team, which means I don't just demo features. I work directly with Databricks customers to enable platform features and run them in a production environment with real users, real traffic, and real workloads. My main focus at Databricks is machine learning, generative AI, and agent-based applications, such as Databricks Apps, Databricks Model Serving, and most recently, Lake Lakehouse.
And I want... Can you hear me? Good. I want to make it clear where this material comes from. Everything I'm going to show you today is not taken from the board. This is the result of implementations with customers. Some went smoothly, some painfully, including a deployment at a critical energy agency that manages the nation's complex power grid. So, when I say " production patterns," I mean approaches designed for the realities of the live environment.
Bringing these systems to production requires a scalable architecture capable of handling concurrent users, multiple AI agents, and sustained transaction volumes. Before I get into the details, I want to take 30 seconds to assess the audience experience. Raise your hand if you have ever built or deployed a Databricks App. Perfectly. What about an application that connects to a database in production? Good. OK. So, many of you are already on the front lines.
Now raise your hand again if you've encountered a mysterious authentication error that took 30 minutes to resolve. Okay. Good. I see those hands and I know exactly what mistake you're thinking about. This mistake and the four patterns that prevent it are what we'll be talking about in the next few minutes. So, let me outline the situation. You've built a great application on Databricks. The demo was successful. Stakeholders are delighted .

01:58 The 61st minute authentication error

You release it, and for the first hour everything works perfectly. Requests are coming in. Your application is responding. Users are satisfied. But then , at the 61st minute, not an hour, not two, exactly 60 minutes after the cold start of the application, you see the error on the right of this slide. Truth? Your intuition: perhaps a mistake in secrets. The service principal may have been revoked. There may be a bug in the SDK .
You spend the next 30-60 minutes searching for credentials, redeploying , opening tickets, and then you realize that your OAuth token just expired. Lakehouse uses OAuth tokens as passwords to Postgres. These tokens have a strict 60-minute expiration limit, and Postgres only checks the expiration date at login. Open connections continue to work, but as soon as a new connection is tried to open after the 60th minute , it throws an error.
If you didn't know about this behavior in advance, you just wasted an hour, maybe more. And here's the trickiest part. This is not just one mistake. These are four anti-patterns superimposed on each other. So, on the right you see that we don't have proactive token rotation. We assume that the built-in PG bouncer handles OAuth, but this is not the case. And each query to the database actually looks for threat metadata.
And this puts a very heavy load on our database. And lastly: parallel agent entries using primitive insert. So, we have one mistake, four root causes, and today we will fix all four. Okay, before we get to the patterns, one slide with important context, as the situation has changed recently, as of March of this year. Lakehouse is a fully managed serverless Postgres from Databricks. Postgres 17 with PG vector, PG stat statements, and all the tools you already know, like PSQL, PG admin, etc.
It became generally available in February 2026 on AWS, and soon on Azure. There are two modes: provision—this is an outdated model, and autoscaling—a more modern option. As of March of this year, new instances of the provision type can no longer be created. Everything runs on autoscaling, and everything in this presentation assumes it: Postgres 17, scaling to zero, instant copying, database branching, and the Postgres project API .
By the way, existing provision instances are automatically upgraded to this more modern version this month, but you also have a migration path with PG dump and PG restore. So, here's our plan for the remaining minutes. We will look at four patterns that together form a complete, production-ready Lakehouse integration within a single Databricks application built on FastAPI and Streamlit—a very popular stack. Think of these patterns as layers of a defense system.

05:55 Four production patterns overview

Patterns one and two are about authentication. We'll make sure that the connection pool never uses an expired token, not once, at all. Pattern three is about reading so that critical paths don't overload your database. Pattern four concerns recording so that parallel agent calls do not corrupt each other's state. All four patterns connect to Lakehouse autoscaling via the standard Postgres protocol. And the beauty is that it's just Postgres, but the OAuth model is a little different, and that's what confuses a lot of developers .
Okay, let's start with the basics. So, OAuth and a self- managed pool. Lakehouse supports two authentication methods. The first is authentication using an OAuth token. It uses your Databricks ID. The token is passed as a Postgres password. This is the right choice for Databricks applications and service principles because the platform manages identity for you. But there is a caveat: tokens expire exactly in 60 minutes, and verification is performed upon login.
The second method is self- authentication using the Postgres password. A static role with a password, without rotation, is best suited for long-running services, third-party tools, or something that physically cannot change credentials every hour. For agent applications running inside Databricks, you should choose OAuth. You get the Databricks identity model, management via Unity Catalog, and deep integration with the SDK.
But, and this is critical, if you choose OAuth, you take on the problem of rotation. The platform won't do this for you, which brings us to the OAuth in the pool trap. There are two facts that need to be learned. Fact one: OAuth tokens are valid for 1 hour. Attempts to reconnect after the expiration date fail. Open connections continue to work. Fact two: Lakehouse's built-in PG Bouncer connection pooler does not support OAuth.

07:48 OAuth token expiry and the connection pool trap

It only works with Postgres' own password-protected roles. Putting this together, we have: there is no pooler, so you need to manage your own connection pool. However, we have the appropriate tools. So, we have an asynchronous psycopg3 connection pool, but we still have two types of failures to protect against. The first is stale tokens on new connections, and the second is pool exhaustion when all slots are occupied and requests start to time out .
The first pattern solves the first problem, and the second one solves what the first one cannot handle. Let me show you both. The first pattern is elegant in its simplicity. You create a subclass, essentially a subclass of psycopg. connection, call it, say, OAuth connection, and override the connect method of the class. Whenever psycopg3 opens a new physical connection, it calls your subclass. You use the Databricks SDK to get a fresh token at this point, enter it as a password, and call the parent connect.

09:01 Custom psycopg3 connection class for fresh tokens

As you can see on the right. You pass this class to the asynchronous connection pool as a connection class. Now, each new physical connection at startup, after displacement, or scaling carries a fresh token. What this doesn't solve is already open connections. If your pool opened the connection at the very beginning, they are still using the initial token. So, after the 61st minute, these particular connections remain open and working, but if the pool tries to crowd out and reopen one of them, the new connection will use OAuth, get a fresh token, and work fine.
The problem is what to do with the connections that were active from the beginning and have not yet been displaced? This is where the second template comes in handy . The second pattern is a background thread. This is a "pool watchdog" that runs for the entire lifetime of your application. That's what he does. At startup, your application starts and pool A is created with fresh tokens. The background thread reads the expiration time from the credential response and goes to sleep until " expiration time minus 300 seconds".

10:13 Pool watchdog: recreate connection pools before tokens expire

Approximately 5 minutes before the token expires. At minute 55, the thread wakes up and creates pool B. It is brand new, with fresh tokens. It automatically and atomically swaps the global pool reference so that new requests go to pool B. It then calls a shutdown on pool A with the " wait equals true" parameter, allowing current requests to complete before removing pool A. By the time the old token expires at the 60th minute, pool A is already down and pool B is fully serving traffic.
Zero downtime , zero authentication errors. Important warning : As of earlier this year, there is a known SDK buffer error, please be aware of it. The token validation in the SDK does not have a safe buffer and may return a nearly expired token. Therefore, do not trust the expiration time directly from the SDK. Decode it yourself and add a buffer of 60 to 300 seconds, as I show here on the slide. A quick but important point about pool size, as this is where developers make costly mistakes too.
Let's think about productive environments. So, first, the minimum size: keep it at one, two, or three. Leap base automatically scales to zero during downtime. Supporting "warm" connections prevents this. So start small. Maximum size. Choose it according to the peak number of concurrent users, as OLAP bypasses the PG bouncer, and each slot in your pool is a real Postgres server process on the other side. You are not sharing a connection, you are allocating real computing resources.

11:59 Async pool sizing and scale to zero

An 8 CU Leap base instance supports about 1600 maximum connections, and a 16 CU instance supports about 3300. So the real bottleneck is the parallelism of your application, not the limitations of Postgres. And remember about hard limits. Leap base has a 24-hour inactivity timeout and a three-day maximum connection lifetime. Good. Now that your pool is reliable, let's talk about reading. In agent applications, the busiest read path is almost always the search for stream metadata.
Each incoming request contains a thread ID, and the application needs to retrieve the thread metadata to restore the context before invoking the agent. At 100 requests per minute, if each request accesses the database, you have 100 DB requests per minute on a path that shouldn't be that busy. The flow metadata does not change at each step. The solution is a cache of least recently used items with a time to live, i. e. an LRU + TTL cache .

13:18 LRU cache with TTL for hot thread lookups

A TTL of 5 minutes with 256 cache keys means that at 100 requests per minute you go from 100 requests to the database every 5 minutes to about three. You break this loaded path for the vast majority of requests. And you have two implementation options. You can use the time LRU cache decorator, which wraps the LRU cache, or the TTL cache as a call parameter. Both work. Key design decisions: A TTL of three to five minutes balances data freshness with workload.
Maximum size from 256 to 1024 to cover sessions and always clear the cache after writing with cache clear. There is one catch here. Cache clear actually deletes the entire cache. Therefore, if you need to selectively delete data by key, you can use cachetools. ttlcache, which supports this. The last template concerns records. In particular, the race conditions that every agent application eventually faces. Here is the scenario.

14:34 Upserts and race conditions with concurrent agent writes

The agent processes the user step. At the beginning of a conversation, a user message is received and the application attempts to record it. We have an insert for thread 42, step five, and of course any message that needs to be saved. After 4 ms, the result of the tool call returns, and the parallel process attempts to write the tool results to the same step. So, again we have thread 42, step five plus a message from the tool.
And we have a conflict between thread 42 and step five; with a simple insert one of these records will fail. With a properly constructed upsert, the newer entry wins and data integrity is preserved. The SQL pattern is insert on do update . Just like in the example on the slide. The where condition ensures that only the newer entry wins. If an older record arrives late for some reason, it will not overwrite the fresh data.
There are two options. You have do update, plus excluded, plus a where clause for updating in order, like agent queues, plus do nothing and an idempotence key for repeat events where you just want to guarantee delivery at least once without duplication. Critical nuances of the scheme. On conflict requires a pre- existing unique constraint or index on the thread ID, in this case an internal index . If there is no constraint, your unit test will pass because there is no concurrency, but you will get a runtime error on the first concurrent write in production.
So define the limitations explicitly before deployment. Now let's get back to the real world. The system that prompted all four of these patterns is an intelligent assistant for a federal agency's infrastructure application. This organization is responsible for managing the national energy distribution. The architecture is a single application of Databricks, FastAPI, and Streamlit running on top of two separate data layers in Lakehouse.
We have a transactional layer that is recorded by the application in real time and contains three tables. Conversation flows, agent queues, and tool events. This is where all four patterns apply. Each agent record goes through templates for updating (upsert). Each read for the thread context goes through the third template. The pool running below uses patterns one and two to stay up to date. A synchronized read-only layer powered by Lake Flow with Unity Catalog contains an object directory, equipment classification, and regulatory documents .

17:05 Real-world energy knowledge assistant architecture

These are the knowledge base tables that the agent accesses. Nightly synchronization processes the bulk of the knowledge base. Continuous synchronization at intervals of approximately 15 seconds keeps priority objects up to date. So we have one application, two data templates, and four persistence templates that tie it all together. Our time is almost up, and here are four templates for you to take with you today. The first is the token life cycle.
Update the pool proactively. The second is pool ownership. PG Bouncer is not an option, but we have an asynchronous connection pool cycle PG 3 plus OF connection. And remember, adjust the size to fit your users, not Postgres limits. Number three— hot reads. You can cache stream metadata. This will reduce the load on the database by orders of magnitude. Remember, you can also use cachetools if you need to perform invalidation on individual keys.
And lastly—the correctness of upsert operations. Define uniqueness explicitly. Remember excluded plus where for in- order updates, and do nothing plus the idempotence key for retries, and check for restrictions before going into production. This is the main thing I want you to take away from today's meeting. So, the laptop is working. Now you know how to bring this into production for real users. Thank you.

Learn more about the Databricks Data and AI platform.