Non parallel plan only in Azure when SESSION_CONTEXT is used

Question

I'm observing strange difference between query plans that I'm getting on my local machine and on Azure SQL. I'm trying to implement row level security, where I read user identifier from SESSION_CONTEXT and then in TVF I check whether the user has access.

On my local machine - SQL Server 2019 Developer edition, DB in compatibility level 150 the query plans are as expected. But when I run it on Azure DB which is also 150 compatibility level, I only get non-parallel query plans with NonParallelPlanReason="NonParallelizableIntrinsicFunction" . I tried a Hyperscale database as well as a DB that is in Elastic Pool and the result is same on both DBs.

You can reproduce that with following code:

CREATE TABLE Users (
    UserIdentifier nvarchar(100) PRIMARY KEY CLUSTERED
)

INSERT INTO Users (UserIdentifier) VALUES ('MyUserIdentifier')

CREATE TABLE TableWithRLS (
    Id int NOT NULL IDENTITY(1,1) PRIMARY KEY CLUSTERED,
    DataColumn nvarchar(100) NULL
)

INSERT INTO TableWithRLS (DataColumn) 
SELECT TOP 10000000 A.[name] FROM sys.all_columns  AS A
CROSS JOIN sys.all_columns  AS B
CROSS JOIN sys.all_columns  AS C


CREATE OR ALTER FUNCTION CheckAccess (@userIdentifier varchar(100))
RETURNS TABLE
WITH SCHEMABINDING
AS
RETURN
    SELECT TOP 1 1 AS HasAccess FROM dbo.Users WHERE UserIdentifier = @userIdentifier

EXEC sp_set_session_context N'UserIdentifier', N'MyUserIdentifier', 1

-- This query gets always non-parallel query plan on Azure
SELECT MAX(DataColumn) FROM TableWithRLS AS X
CROSS APPLY CheckAccess(CAST(SESSION_CONTEXT(N'UserIdentifier') AS nvarchar(100)))

When I select the value from session context into a variable first, then it generates parallelizable query plan even in Azure.

DECLARE @userIdentifier AS nvarchar(100) = CAST(SESSION_CONTEXT(N'UserIdentifier') AS nvarchar(100))
SELECT MAX(DataColumn) FROM TableWithRLS AS X
CROSS APPLY CheckAccess(@userIdentifier)

Unfortunately I can't do that (or at least I'm not aware how to do that) because I need an inline TVF.

Query plan from Azure: https://www.brentozar.com/pastetheplan/?id=ByxZm45e9

Query plan from local: https://www.brentozar.com/pastetheplan/?id=BylHXV9lc

Is there any difference in SESSION_CONTEXT implementation in Azure that could be causing that? Or does any one has any other ideas what could be the issue?

Answer 1

Is there any difference in SESSION_CONTEXT implementation in Azure that could be causing that?

Yes. Although Azure SQL Database and SQL Server are built from a common code base, it is not at all unusual for one to be ahead of the other. Despite Microsoft marketing in years past, it is not always the Azure edition that is ahead.

Disallowing parallel plans when SESSION_CONTEXT is used is a feature flag (DisableSessionContextParallelPlan) that is currently enabled in Azure, without any way to turn it off. It can be enabled on SQL Server with undocumented trace flag 11042 at query, session, global, and start-up level.

Using a Stack Overflow demo database (any parallel query will do):

EXECUTE sys.sp_set_session_context 
    @key = N'key', 
    @value = 123,
    @read_only = 1;

SELECT COUNT_BIG(*) 
FROM dbo.Badges AS B 
WHERE B.UserId = CONVERT(integer, SESSION_CONTEXT(N'Key'))
OPTION (QUERYTRACEON 11042);

The query is parallel without the trace flag on SQL Server.

These options are often added to preview features, as a response to rare bug conditions, or where a security risk exists.

A bug fix for incorrect results when using SESSION_CONTEXT in parallel plans was released in SQL Server 2019 CU14.

---

There were problems with the fix, reported in the documentation for CU16:

SQL Server 2019 CU14 introduced a fix to address wrong results in parallel plans returned by built-in SESSION_CONTEXT. However, this fix could cause Access Violation dump files while resetting the SESSION for reuse. To mitigate this issue, you can disable the original fix, and disable the parallelism for built-in SESSION_CONTEXT to avoid wrong results. To do this, use the following trace flags:

11042 - This trace flag disables the parallelism for built-in SESSION_CONTEXT.

9432 - This trace flag disables the fix introduced in SQL Server 2019 CU14.

Microsoft is working on a fix for this issue that will be available in a future CU.

Answer 2

From the XML of the AzureDB plan: NonParallelPlanReason="NonParallelizableIntrinsicFunction"

This is due to the context intrinsic being used and Azure SQL DB specifically setup to disable parallelism on these intrinsics (but this is not the case for box unless specific items are enabled). There are ways to make the box product (on-prem) work the same, but I believe you'd like it the other way around.

Essentially, this is how Azure SQL DB is currently configured, and I don't know if there are ways to get around this in your subscription (I don't generally work with Azure SQL DB). I couldn't find anything in the current documentation to describe this behavior.