Connecting Azure Databricks to AI Gateway in API Management with Modern Authentication

Gateways all the way down

If you thought we were done around the interoperability stories that exist between Databricks and other Azure services, then I have some news for you. Since the last time we visited this topic, a lot has changed both within Databricks and Azure. Both companies haven’t slowed down their pace of development, doubly-so when it comes to things around agents and AI.

On the Azure side, the biggest leaps have been around Microsoft Foundry (née Azure AI Studio). With a vast selection of models and a suite of agent development tools, it’s an attractive space for enterprises who use Azure to think about for building and hosting agents, or at least using the models it provides as part of a larger project. In addition, there is now the “AI Gateway” functionality that has been added to Foundry, which, surprise: it’s just API Management (APIM) with some hooks into Foundry for automatically configuring routes and services to your projects and deployments. It does more than that, of course, such providing controls around token usage, but the big draw is a central control plane for all your deployed agents.

Meanwhile, Databricks has been doubling down on all things agentic AI as well, and has rebranded and enhanced legacy services into what it calls Unity Gateway (née Mosaic AI Gateway, among others) which provides features automatic model and harness routing, tracing and observability, and all the strong governance controls of Unity Catalog applied to the agentic world: agents, models, model providers, MCP servers, and more.

While some people might look at the different capabilities and think this is an “either/or” scenario, the reality is that most Azure customers are going use (or even need to use) both: developers building apps and agents in Databricks, for example, may want to leverage a model or other agent deployed in Foundry. Meanwhile, developers in Foundry may need to call a Databricks managed MCP service.

These types of calls aren’t all that difficult; in fact, both platforms provide built-in connectors or UI solutions to make this happen: simply provide the service endpoint, pick your authentication method, and you’re pretty much done. However, the introduction of the respective gateways, especially on the Foundry side, can make this just a little more complicated. Let’s look at one example.

Modern problems require modern (authentication) solutions

Here’s the scenario: you’re a Databricks developer who is building a new agent, and you want to use a model hosted by AI Foundry. Seems simple enough: either you (or your platform/admin team) will take care of setting up a new model provider service in Unity Gateway. In the UI, it’s even directly called out:

You’ll pick which catalog and schema the provider object will live in, then you give explicit permissions to the provider to your users/groups/service principals that you want to be able to use it to call the models it provides, or, create explicit model objects that call one of the models that the provider, uh, provides (or, if your team wants to limit the models it provides, they can be selective about what models to allow from the provider).

Then, you just provide the URL to your foundry resource, and you’re… pretty much done. Except for that one little section labeled “Auth method.” This is where you’ll decide how, when the call is made from Databricks to Foundry, how does it authenticate to the service? And you’ve got three choices here, and frankly, only one right one: Service Credentials. These are managed identities in Azure that can be invoked by Databricks to generate access tokens for specific audiences that then get passed along with the call to the service, without the need for any API keys, service credential keys, or other things you have to worry about protecting.

Outside of then granting the Databrick Access Connector that is tied to that service credential to appropriate roles on the Foundry service in the portal, that’s pretty much it.

However, organizations that want to take advantage of the AI Gateway feature of APIM and Foundry, the math changes a little bit. Now, your Databricks provider isn’t authenticating to Foundry anymore, it’s authenticating to APIM. And APIM, out of the box, doesn’t follow the same rules of authentication that going to Foundry directly provides. Instead, you’d normally pick the “Custom” option, but when you do, the Auth methods change:

As you can see, the service credential option is gone, and what’s more, you have to manually provide the model names, and the capabilities of each one (such as chat completions, responses, etc.). Then you really only have API key or custom header authentication. In other words, no modern authentication solution out of the box. Bummer.

Where there’s smoke, there’s fire

This was the exact situation I found myself in when supporting a customer recently: their requirement was no key exchanges, and instead, relying on modern authentication for all of this to work. All the documentation is pretty clear: if you use a custom provider, you’re stuck with key exchanges only. Clearly though, when you pick the service credential option, Databricks is doing something to request an access token on the Azure side and passing that along to Foundry for authentication. So that got me thinking: what if we still picked the Foundry option for the provider and gave it the APIM endpoint. What would happen?

Well at first, exactly what you expect: for the purposes of this blog, understand that this was an out-of-the-box deployment of Foundry and AI Gateway; by default, the only thing that the APIM expects is a subscription key as the authentication mechanism. However… remember that AI Gateway is just APIM: we can do whatever the hell we want to the inbound policies!

So consider this: if we know what the audience URL is what Databricks creates the token for, we can craft an inbound authentication policy on the API side to accept it. Well, I’ll save you a lot of time and say that after doing a lot of tracing and logging, it’s the audience https://cognitiveservices.azure.com/.default.

Then, if we go to API management and find our matching Foundry resource, and look at the policy for all of the operations, all we need to do is add a validate-azure-ad-token (docs: https://learn.microsoft.com/en-us/azure/api-management/validate-azure-ad-token-policy):

<validate-azure-ad-token tenant-id="<YOUR AZURE TENANT ID>" header-name="Authorization" failed-validation-httpcode="401" failed-validation-error-message="Unauthorized. Access token is missing or invalid." output-token-variable-name="jwt">
    <client-application-ids>
        <application-id><YOUR CONNECTOR CLIENT ID HERE></application-id>
    </client-application-ids>
    <audiences>
        <audience>https://cognitiveservices.azure.com</audience>
        <audience>https://cognitiveservices.azure.com/</audience>
    </audiences>
</validate-azure-ad-token>

You’ll just append this to your inbound policy for “all operations” and now, you’ve secured your API by only allowing identities with matching client IDs and the required claims for Foundry to use the APIs. How does it work? Well, first, in the main policy definition, we’re stating that we only want to allow tokens generated from our tenant to be allowed, then we’re filtering on a specific list of client IDs of the identities, and we’re specifying that the token must have been generated for the same audience as Foundry.

So now, back in Databricks, we can configure the priovider to use Foundry, and provide the base URL to our APIM endpoint:

Again, depending how you deployed your AI Gateway for Foundry, you just need to make sure your Foundry resource URL is of the format https://<your apim resource name>.azure-api.net/<your foundry resource name>

Then, you should be able to test it. In the bottom of the provider window, using the cURL method, try calling it but make sure you select an actual model you’ve deployed in Foundry, or you’ll get a 404:

Heads-up: Look, this isn’t documented functionality. It’s reversed engineered from how Databricks talks to Foundry natively. There’s no magic here, but this is a workaround, so treat it as such. Hopefully Databricks will add “Microsoft AI Gateway” as an option natively. Until then, this works.

Some real talk:

In no particular order, here’s some other things to keep in mind when you’re setting out to get this working. Bear in mind, part of my day job is helping customers implement things like this successfully, and if there’s one thing you know about me, is that your boy has opinions:

• First and foremost, this example is based on an out-of-the-box setup for AI Gateway in Foundry; if your organization has already customized the APIM policies for other reasons, don’t just rip and replace; work with your team to implement the policy so that it doesn’t break your other apps, services, or users. There’s control flow logic in APIM, but that could be a whole separate blog post.
• Given where I work, I’d tell you Unity Gateway is pretty darn good at managing (and governing) your AI estate, and you should give it serious consideration. Even if you want APIM to be your enterprise solution for all things APIs, not just AI-related APIs in your organization, Unity Gateway is still a table-stakes piece of functionality for people working Databricks. Use it.
• Building on that: realize that APIM for Foundry is an additional cost to you and your organization. And the requirements for AI Gateway-enabled APIM is pretty steep. Read the fine print and requirements carefully, especially if you need private endpoint support, VNet integration, more capacity for requests, and more. Meanwhile, Unity Gateway is a free piece of functionality in Databricks.
• APIM mastery is a skill in itself; policy definitions and management can be hard for newcomers. Use whatever tools make sense to help you build effective, safe, and effective policies. And don’t test it in your production APIM instances! You do have a dev instance… right?

Hope this helps! Remember: Azure Databricks is part of your Azure ecosystem and a first party service on Azure. So, use it as such, especially when it comes to modern authentication.