Why use APIM with Data API?
The Data API is one of the underrated features of Lakebase. Traditionally, teams that wanted to share a database over HTTP built a layer of REST services in front of it. Then they had to keep that layer alive: endpoints for each table or query they wanted to share, authentication, input validation, deployments, and someone on call when it broke. The Data API takes most of that work off their plate. Once it is turned on for a database, every table and function that a role's grants allow is reachable over REST, with filtering, paging and writes in the PostgREST style. Postgres grants decide what each caller can touch.
That serves applications that already work with Databricks well, such as a Databricks App, a job, or a service that holds a Databricks service principal. Enterprises, however, rarely share data only with themselves. A distributor wants to place orders from its own purchasing system. A chain of cafés that buys from the supplier wants its own systems to check the status of its orders. Another business unit, whose applications already sign in through Microsoft Entra ID, wants to read the product catalog. Each of these clients has its own identity and its own contract. Automatic identity management can bring service principals from the organization's Entra ID directory into Databricks, so identity is less of a hurdle than it used to be. Each application would still have to learn to obtain Databricks tokens, though, and each would need its own Postgres role and grants.
The platform team then faces a different set of questions:
- How do we let every client keep using the Entra token it already has, instead of teaching each application to obtain Databricks tokens?
- How do we stop one partner's runaway script from slowing the database down for everyone else?
- How do we make sure each client sees only the data it is entitled to?
- When a contract ends, how quickly can we cut that client off?
Many organizations on Azure already answer these questions with Azure API Management (APIM), the gateway they use to publish their other APIs. APIM onboards clients through products and subscriptions and authenticates them with the identities they already have in Entra ID. It can also enforce rate limits and quotas per client, so a runaway script gets 429 Too Many Requests from APIM before its calls reach Lakebase. Placed in front of the Data API, it lets each side keep its own job: APIM governs which clients call the API and how often, while Lakebase and Postgres govern what each request can read or change.
APIM is one extra component to run and pay for, so it is worth checking whether you need it. If the only clients are Databricks Apps or internal services that already hold Databricks credentials, calling the Data API directly is simpler. The case for APIM grows with the number of clients from outside Databricks. It is strongest when those clients belong to different organizations, such as partners or customers, whose data must stay apart.
Figure 0. Where Azure API Management fits with the Data API
In this blog, we will walk through the two pieces that make this work:
- APIM needs a Databricks token to call the Data API.
- Each client must reach Postgres with an identity that says which organization it belongs to, so that row-level security (RLS) can restrict what it can access.
The journey of a client's request
Let's start with a few terms. A client is an application that calls the API, such as the distributor's purchasing system from earlier. The API owner decides which data each client can reach, and that decision travels in the client's token, so a client can't choose its own access. Several clients can share the same access, for example the café chain's ordering system and its reporting job, and they all see the café chain's orders and nothing else. The paragraphs below introduce the Azure pieces that carry this decision from the client's token to the database.
On the Azure side, identities live in Microsoft Entra ID. An app registration is how an application gets an identity there, and this design uses two kinds. Each client has its own app registration, which it uses to sign in. The API in APIM has one too, which we call the gateway app registration. There is only one, however many clients call the API. It represents that API in Entra ID: clients ask Entra ID for a token for the gateway app registration, and APIM accepts only tokens issued for it. Several APIs in APIM can share one app registration, so that a single token works for all of them, but in this blog it represents one API.
An app role is a named permission that belongs to the API. In this design, each app role decides whose orders a client can reach. The API owner defines the app roles Orders.AppRoleA and Orders.AppRoleB on the gateway app registration and assigns one of them to each client's app registration, so clients with the same app role share the same orders. The client takes no part in this. It asks Entra ID for a token for the gateway, and Entra ID adds the app role the API owner assigned to it. If the API owner hasn't assigned one, Entra ID refuses to issue the token.
Figure 1. The API owner assigns each client one app role. Clients with the same app role see the same rows and nothing else.
A managed identity is an identity that Azure creates and manages for a resource such as APIM. Azure hands out tokens for it on request, so there is no password or secret for anyone to store. In this design, the API owner gives APIM one managed identity per app role, and each managed identity gets its own role in Postgres. That role is the only member of a Postgres group role named after the app role, such as orders_app_role_a. (A group role is a Postgres role that other roles join, and its members inherit its grants.) That is how the app role travels from the client's token all the way to the database.
The animation below puts these pieces together, one step at a time:
Figure 2. Clients, app roles, identities and rows, step by step. It loops, starting again from the title card.
With those terms in place, there are two identities at work on every request. The client signs in to Entra ID with its application ID and a credential (a client secret, a certificate or a federated credential) and sends the Entra token it receives to APIM. APIM checks that token, reads its app role, and then calls Lakebase with a different identity: the managed identity of that app role. The client's token stops at APIM and never reaches Lakebase.
Getting a Databricks token for a managed identity takes APIM two extra calls, one to Entra ID and one to Databricks, which adds a few hundred milliseconds to a request. A Databricks token stays valid for an hour, so APIM keeps one per managed identity in its cache and reuses it until five minutes before it expires. These two calls therefore happen only on the first request for each app role, and then about once an hour.
Getting a Databricks token with no secrets
The Data API accepts tokens issued by the Databricks workspace, so APIM needs a way to turn its managed identity into a Databricks identity. Databricks workload identity federation does exactly that. We register the managed identity as a service principal in the Databricks account and attach a federation policy to it. A federation policy is a trust rule: it names an identity provider (the issuer), the audience a token must be issued for, and the identity it must represent (the subject). When APIM presents a matching Entra token at the workspace's `/oidc/v1/token` endpoint, Databricks returns a Databricks token for that service principal, so no secret is stored anywhere.
Below are the steps:
- Create a user-assigned managed identity and attach it to APIM
- Register it in Databricks with a federation policy
- Give it a Postgres role and grants
- Add the token exchange to the APIM policy
Create a user-assigned managed identity and attach it to APIM
We use a user-assigned managed identity rather than APIM's built-in one, because it can be created, registered in Databricks and granted in Postgres before APIM exists. The identity has two IDs that matter here: its client ID, which becomes its Databricks application ID and its Postgres role name, and its object ID, which appears as the sub claim of its Entra tokens.
Register it in Databricks with a federation policy
The managed identity becomes an Entra-managed service principal in the Databricks account, with the client ID as its application ID, and gets assigned to the workspace. The federation policy then tells Databricks which Entra tokens may be exchanged for it. In Terraform, it looks like this:
locals {
token_exchange_app_id = "<microsoft's-fixed-id-for-azure-ad-token-exchange-api>"
}
resource "databricks_service_principal_federation_policy" "apim" {
provider = databricks.account
service_principal_id = tonumber(databricks_service_principal.apim.id)
policy_id = "entra-apim-mi"
oidc_policy = {
issuer = "https://login.microsoftonline.com/${var.tenant_id}/v2.0"
audiences = [local.token_exchange_app_id, "api://AzureADTokenExchange"]
subject = azurerm_user_assigned_identity.apim.principal_id
}
}Note the usage of two audiences: the policy must accept both the URI and its ID. The subject is the managed identity's object ID.
Give it a Postgres role and grants
Inside Postgres, the Data API connects as a role called authenticator and switches to the role named in the token. The managed identity's role has to be created with the `databricks_create_role` function and granted to authenticator. After that, ordinary grants decide what it can do:
CREATE EXTENSION IF NOT EXISTS databricks_auth;
SELECT databricks_create_role('<mi-client-id>', 'SERVICE_PRINCIPAL');
GRANT "<mi-client-id>" TO authenticator;
GRANT USAGE ON SCHEMA public TO "<mi-client-id>";
GRANT SELECT, INSERT, UPDATE, DELETE ON public.products TO "<mi-client-id>";
Add the token exchange to the APIM policy
The policy below runs on every request. It looks for a cached Databricks token, and on a miss it gets an Entra token for the managed identity, exchanges it, and caches the result until five minutes before it expires. We have trimmed error handling from the excerpt, and the full policy is in the repository.
<cache-lookup-value key="lakebase-dbx-token" variable-name="dbx-token" />
<choose>
<when condition="@(!context.Variables.ContainsKey("dbx-token"))">
<authentication-managed-identity resource="api://AzureADTokenExchange"
client-id="{{lakebase-mi-client-id}}" output-token-variable-name="entra-token" />
<send-request mode="new" response-variable-name="exchange-response" timeout="20">
<set-url>{{lakebase-workspace-url}}/oidc/v1/token</set-url>
<set-method>POST</set-method>
<set-header name="Content-Type" exists-action="override">
<value>application/x-www-form-urlencoded</value>
</set-header>
<set-body>@("grant_type=" + System.Net.WebUtility.UrlEncode("urn:ietf:params:oauth:grant-type:token-exchange")
+ "&subject_token=" + (string)context.Variables["entra-token"]
+ "&subject_token_type=" + System.Net.WebUtility.UrlEncode("urn:ietf:params:oauth:token-type:jwt")
+ "&client_id={{lakebase-mi-client-id}}&scope=all-apis")</set-body>
</send-request>
<set-variable name="exchange-json" value="@(((IResponse)context.Variables["exchange-response"]).Body.As<JObject>())" />
<set-variable name="dbx-token" value="@((string)((JObject)context.Variables["exchange-json"])["access_token"])" />
<cache-store-value key="lakebase-dbx-token" value="@((string)context.Variables["dbx-token"])"
duration="@(Math.Max(60, (((JObject)context.Variables["exchange-json"]).Value<int?>("expires_in") ?? 3600) - 300))" />
</when>
</choose>
<set-header name="Authorization" exists-action="override">
<value>@("Bearer " + (string)context.Variables["dbx-token"])</value>
</set-header>
<set-backend-service base-url="{{lakebase-data-api-url}}" />
The values in double braces are APIM named values, which are APIM's settings store, and none of them is a secret: the workspace URL, the Data API URL and the managed identity's client ID. The caller's own Authorization header and subscription key are removed earlier in the policy, so only the Databricks token reaches Lakebase.
The alternative: a Databricks OAuth secret
If a managed identity isn't available, APIM can use a Databricks-managed service principal with an OAuth secret instead. The policy calls the same /oidc/v1/token endpoint with grant_type=client_credentials and the secret, and caches the token the same way. There is no managed identity or federation policy to set up, and the Postgres side is identical, with the service principal's application ID as the role name.
The cost is a secret to store and rotate. A Databricks OAuth secret lasts up to 730 days and a service principal can hold five, so rotation works by overlap: create the new secret, update the APIM named value, then delete the old one. Keep the secret in Azure Key Vault and reference it from a secret named value. Both variants behave the same apart from how the token is obtained, so we recommend the managed identity unless your setup rules it out.
One managed identity per app role
So far every request reaches Postgres as the same role. To keep each client's orders apart, Postgres needs to know which app role is calling. The Data API has no way to pass the caller's identity into SQL (no custom settings per request, no pre-request function), so the Postgres role itself has to carry the app role. That means one managed identity per app role, and APIM picking the right one for each request.
Let's use a small example to make this concrete. A supplier publishes one product catalog, and each of its customers places orders and sees only its own. The catalog is shared and read-only for clients, while each order records the group role it belongs to:
|
Schema |
Objects |
Controlled by |
|
public |
products, the shared catalog |
Grants: clients can read it, not change it |
|
app |
orders, order_lines and the view order_totals |
RLS: each group role reads and writes only its own rows |
|
private |
The function owner_role(), which returns the caller's group role |
Not exposed through the Data API |
Below are the steps:
- Give the gateway app registration its app roles
- Create a managed identity per app role
- Pick the managed identity from the caller's app role
- Add a group role per app role and turn on RLS
Give the gateway app registration its app roles
The gateway app registration represents our API in Entra ID, and it defines the app roles Orders.AppRoleA and Orders.AppRoleB. Each client's app registration is assigned exactly one of them. We also set "assignment required" on the gateway, so Entra refuses to issue gateway tokens to applications that hold no role. A client then asks Entra for a token with the scope api://<gateway-app-id>/.default, and the token comes back with roles: ["Orders.AppRoleA"].
The client apps in this blog authenticate with a client secret, which keeps the curl examples short. Production clients should use a certificate, a federated credential or, if they run in Azure, their own managed identity.
Create a managed identity per app role
Each app role gets its own user-assigned managed identity, attached to APIM, registered in Databricks with its own federation policy, and given its own Postgres role. These are the same steps as in the previous section, repeated per app role.
Pick the managed identity from the caller's app role
The policy for this API validates the caller's Entra token, reads its one Orders.AppRole<X> app role, and uses it to choose the managed identity and the cache key. Requests without exactly one such app role are refused before anything else happens:
<validate-azure-ad-token tenant-id="{{entra-tenant-id}}" output-token-variable-name="client-jwt">
<audiences>
<audience>api://{{lakebase-gateway-app-id}}</audience>
</audiences>
</validate-azure-ad-token>
<set-variable name="app-role" value="@{
var jwt = (Jwt)context.Variables["client-jwt"];
var roles = jwt.Claims.ContainsKey("roles") ? jwt.Claims["roles"] : new string[0];
var matches = roles.Where(r => r.StartsWith("Orders.AppRole")).ToArray();
return matches.Length == 1 ? matches[0] : "";
}" />
<set-variable name="app-role-key" value="@(((string)context.Variables["app-role"]).Replace("Orders.AppRole", "").ToLowerInvariant())" />
<!-- ... refuse with 403 if app-role is empty, then look up "lakebase-dbx-token-app-role-" + app-role-key ... -->
<choose>
<when condition="@((string)context.Variables["app-role-key"] == "a")">
<authentication-managed-identity resource="api://AzureADTokenExchange"
client-id="{{lakebase-mi-app-role-a-client-id}}" output-token-variable-name="entra-token" />
</when>
<when condition="@((string)context.Variables["app-role-key"] == "b")">
<authentication-managed-identity resource="api://AzureADTokenExchange"
client-id="{{lakebase-mi-app-role-b-client-id}}" output-token-variable-name="entra-token" />
</when>
</choose>Note the usage of one <when> branch per app role. Each branch names its app role's managed identity through a named value, so adding an app role means a new branch and a new named value, along with the other pieces per app role.
Add a group role per app role and turn on RLS
On the Postgres side, each app role gets a group role that never logs in, with the managed identity's role as its only member. The grants go to the group role, and the member inherits them. The function private.owner_role() returns the caller's group role from pg_auth_members, a catalog that every role can read. Each table under RLS carries an owner_role column that defaults to the caller's group role, and one policy per table filters reads and checks writes:
CREATE ROLE orders_app_role_a NOLOGIN;
GRANT orders_app_role_a TO "<app-role-a-mi-client-id>";
CREATE FUNCTION private.owner_role() RETURNS text
LANGUAGE sql STABLE SET search_path = pg_catalog, pg_temp
AS $$
SELECT CASE WHEN count(*) = 1 THEN min(g.rolname) END
FROM pg_auth_members m
JOIN pg_roles g ON g.oid = m.roleid
JOIN pg_roles u ON u.oid = m.member
WHERE u.rolname = current_user
AND g.rolname LIKE 'orders\_app\_role\_%'
$$;
CREATE TABLE app.orders (
owner_role text NOT NULL DEFAULT private.owner_role(),
id bigint GENERATED ALWAYS AS IDENTITY,
status text NOT NULL DEFAULT 'open',
created_at timestamptz NOT NULL DEFAULT now(),
PRIMARY KEY (owner_role, id)
);
ALTER TABLE app.orders ENABLE ROW LEVEL SECURITY;
CREATE POLICY owner_role_isolation ON app.orders
USING (owner_role = (SELECT private.owner_role()))
WITH CHECK (owner_role = (SELECT private.owner_role()));
GRANT SELECT, INSERT, UPDATE, DELETE ON app.orders TO orders_app_role_a, orders_app_role_b;
The figure below follows a request from a client with Orders.AppRoleA through the lookup:
Figure 3. private.owner_role() finds the group role that the request's role belongs to, and the policy keeps only that group role's rows.
Note the usage of (SELECT ...) around the lookup. It lets Postgres look up the group role once per query instead of once per row. If a role belongs to more than one group role, the function returns nothing, so that caller sees no rows and can't write any. A role in no group role holds no grants at all, so Postgres refuses it with permission denied. Order lines follow the same pattern, and they point at their order through the composite key (owner_role, order_id), for a reason we cover in the gotchas below.
Trying it out
With everything in place, two clients call the same URL with their own tokens: client A has Orders.AppRoleA and client B has Orders.AppRoleB. The animation below plays each kind of request in turn:
Figure 4. Eight requests against the same API, and what each one returns. It loops, starting again from the title card.
The same results as a table:
|
Request |
Result |
|
Client A reads /app/orders |
200, only the orders of orders_app_role_a, with the response header X-Lakebase-App-Role: Orders.AppRoleA |
|
Client B reads /app/orders |
200, only the orders of orders_app_role_b |
|
Client A creates an order |
201, with owner_role filled in as orders_app_role_a by the default |
|
Client A inserts a row with owner_role: orders_app_role_b |
403, 42501, "new row violates row-level security policy" |
|
Client A updates one of client B's orders |
200 with an empty result: client A can't see the row, so nothing changes |
|
Client A adds a line to client B's order |
409, 23503: the composite foreign key finds no order (orders_app_role_a, 3) |
|
Client A changes a product in the catalog |
403, 42501, "permission denied for table products" |
|
No token, or a token for another audience |
401 from APIM, and the request never reaches Lakebase |
Each app role also gets its own cached Databricks token. The first request for an app role shows X-Lakebase-Token-Source: exchange, and the next ones show cache.
Things to watch for
A few details cost us time, or would have, and they are worth knowing before you build this:
- The Data API reports token problems as 400 with a message about the token, not as 401. If your policy evicts a cached token on errors, check the status and the body.
- Missing grants come back as 403 with the Postgres code 42501. So does a role that was never granted to authenticator, with the message "permission denied to set role".
- Create Postgres roles for the Data API with databricks_create_role. Only those can be granted to authenticator.
- The Lakebase project owner and databricks_superuser have BYPASSRLS. RLS never applies to them, and FORCE ROW LEVEL SECURITY doesn't change that.
- Because of that, a view over a table under RLS must be created with security_invoker = true. A view without it runs as its owner, which bypasses RLS, so it returns every client's orders to every client.
- Foreign-key checks ignore RLS. If order lines referenced orders by order_id alone, a client could attach a line to another client's order. Referencing (owner_role, order_id) closes that, because owner_role is always the caller's own.
- When RLS filters out every row, the result is an empty 200, not an error.
Cutting a client off
To offboard a client, remove its app role assignment on the gateway app registration. Entra then refuses new tokens for that client, usually within a few minutes, with AADSTS501051. A token issued before the change keeps working until shortly after it expires, which can be up to an hour later. If you need to block a client immediately, also reject its application ID in the APIM policy.
When this design fits
This works well for tens of app roles. Adding an app role is an infrastructure change (the app role itself, a managed identity with its federation policy, a Postgres role and its group role, and a policy branch), which is easy to drive from Terraform but doesn't scale to thousands. Onboarding another client with an existing app role needs no change at all, only an app role assignment. For thousands of app roles, or isolation per end user, a different design is needed. We discuss the options in the repo's notes on client identity and row-level security.
Related setup
Two questions often come up alongside this design, and both are configured outside it.
- Network access: Azure Databricks can restrict who reaches the workspace with IP access lists and inbound Private Link. If you use them, allow APIM's outbound IP addresses, or connect APIM to a virtual network.
- Partners in their own Entra tenant: a partner's app registration usually lives in the partner's Entra tenant, so it can't hold an app role on your gateway as it is. The partner can make the app multitenant, and an admin in your Entra tenant grants consent to it, which adds it to your directory so you can assign the app role. Alternatively, you create the client app registration in your own Entra tenant, and the partner's workload signs in to it with workload identity federation, so no secret changes hands.
Set it up
Everything in this blog, Terraform and scripts included, is in the lakebase-data-api-apim repo. Its README walks through the setup.
Conclusion
Together, the Data API and Azure API Management turn a Lakebase database into an API you can share widely, without writing or running a service in between. Clients outside Databricks keep the identities they already have, the platform team manages their access in the gateway it already runs, and each client sees only the data its app role allows.
