Database policy support for PUT/PATCH operations - PostgreSQL

config-generators/postgresql-commands.txt

@@ -58,8 +58,9 @@ update Publisher --config "dab-config.PostgreSql.json" --permissions "database_p
 update Publisher --config "dab-config.PostgreSql.json" --permissions "database_policy_tester:create"
 update Publisher --config "dab-config.PostgreSql.json" --permissions "database_policy_tester:update" --policy-database "@item.id ne 1234"
 update Stock --config "dab-config.PostgreSql.json" --permissions "authenticated:create,read,update,delete" --rest commodities --graphql true --relationship stocks_price --target.entity stocks_price --cardinality one
-update Stock --config "dab-config.PostgreSql.json" --permissions "database_policy_tester:create,read"
 update Stock --config "dab-config.PostgreSql.json" --permissions "database_policy_tester:update" --policy-database "@item.pieceid ne 1"
+update Stock --config "dab-config.PostgreSql.json" --permissions "database_policy_tester:create" --policy-database "@item.pieceid ne 6 and @item.piecesAvailable gt 0"
+update Stock --config "dab-config.PostgreSql.json" --permissions "database_policy_tester:read"
 update Stock --config "dab-config.PostgreSql.json" --permissions "test_role_with_noread:create,update,delete"
 update Stock --config "dab-config.PostgreSql.json" --permissions "test_role_with_excluded_fields:create,update,delete"
 update Stock --config "dab-config.PostgreSql.json" --permissions "test_role_with_excluded_fields:read" --fields.exclude "categoryName"

src/Core/Configurations/RuntimeConfigValidator.cs

@@ -47,7 +47,8 @@ public class RuntimeConfigValidator : IConfigValidator
     private static readonly HashSet<DatabaseType> _databaseTypesSupportingCreatePolicy =
     [
         DatabaseType.MSSQL,
-        DatabaseType.DWSQL
+        DatabaseType.DWSQL,
+        DatabaseType.PostgreSQL
     ];
 
     // Error messages for user-delegated authentication configuration.

src/Core/Resolvers/PostgreSqlExecutor.cs

@@ -2,11 +2,13 @@
 // Licensed under the MIT License.
 
 using System.Data.Common;
+using System.Net;
 using Azure.Core;
 using Azure.DataApiBuilder.Config;
 using Azure.DataApiBuilder.Config.ObjectModel;
 using Azure.DataApiBuilder.Core.Configurations;
 using Azure.DataApiBuilder.Core.Models;
+using Azure.DataApiBuilder.Service.Exceptions;
 using Azure.Identity;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Logging;
@@ -146,6 +148,87 @@ private static bool ShouldManagedIdentityAccessBeAttempted(NpgsqlConnectionStrin
             return string.IsNullOrEmpty(builder.Password);
         }
 
+        /// <inheritdoc/>
+        public override async Task<DbResultSet> GetMultipleResultSetsIfAnyAsync(
+            DbDataReader dbDataReader, List<string>? args = null)
+        {
+            // RS1: COUNT of rows matching PK (no policy) — used to distinguish
+            // "row doesn't exist" from "row exists but policy blocked".
+            DbResultSet resultSetWithCountOfRowsWithGivenPk = await ExtractResultSetFromDbDataReaderAsync(dbDataReader);
+            DbResultSetRow? resultSetRowWithCountOfRowsWithGivenPk = resultSetWithCountOfRowsWithGivenPk.Rows.FirstOrDefault();
+            int numOfRecordsWithGivenPK;
+
+            if (resultSetRowWithCountOfRowsWithGivenPk is not null &&
+                resultSetRowWithCountOfRowsWithGivenPk.Columns.TryGetValue(PostgresQueryBuilder.COUNT_ROWS_WITH_GIVEN_PK, out object? rowsWithGivenPK))
+            {
+                // PostgreSQL COUNT(*) returns Int64; convert to int.
+                numOfRecordsWithGivenPK = Convert.ToInt32(rowsWithGivenPK!);
+            }
+            else
+            {
+                throw new DataApiBuilderException(
+                    message: $"Neither insert nor update could be performed.",
+                    statusCode: HttpStatusCode.InternalServerError,
+                    subStatusCode: DataApiBuilderException.SubStatusCodes.UnexpectedError);
+            }
+
+            // RS2: UPDATE result, or UPDATE+INSERT CTE result.
+            DbResultSet dbResultSet = await dbDataReader.NextResultAsync()
+                ? await ExtractResultSetFromDbDataReaderAsync(dbDataReader)
+                : throw new DataApiBuilderException(
+                    message: $"Neither insert nor update could be performed.",
+                    statusCode: HttpStatusCode.InternalServerError,
+                    subStatusCode: DataApiBuilderException.SubStatusCodes.UnexpectedError);
+
+            if (numOfRecordsWithGivenPK == 1) // Row existed — we attempted an UPDATE.
+            {
+                if (dbResultSet.Rows.Count == 0)
+                {
+                    // Row exists but UPDATE returned no rows — update policy blocked it.
+                    throw new DataApiBuilderException(
+                        message: DataApiBuilderException.AUTHORIZATION_FAILURE,
+                        statusCode: HttpStatusCode.Forbidden,
+                        subStatusCode: DataApiBuilderException.SubStatusCodes.DatabasePolicyFailure);
+                }
+            }
+            else if (dbResultSet.Rows.Count == 0)
+            {
+                // Check whether IsFallbackToUpdate was set (passed as args[2]).
+                // If true, the row simply didn't exist — return 404 (same as MsSql's null-RS2 path).
+                // If false (or not present), the INSERT ran but create policy blocked it — return 403.
+                bool isFallbackToUpdate = args is not null && args.Count > 2
+                    && bool.TryParse(args[2], out bool fallback) && fallback;
+
+                if (isFallbackToUpdate)
+                {
+                    if (args is not null && args.Count > 1)
+                    {
+                        string prettyPrintPk = args[0];
+                        string entityName = args[1];
+
+                        throw new DataApiBuilderException(
+                            message: $"Cannot perform INSERT and could not find {entityName} " +
+                            $"with primary key {prettyPrintPk} to perform UPDATE on.",
+                            statusCode: HttpStatusCode.NotFound,
+                            subStatusCode: DataApiBuilderException.SubStatusCodes.ItemNotFound);
+                    }
+
+                    throw new DataApiBuilderException(
+                        message: $"Neither insert nor update could be performed.",
+                        statusCode: HttpStatusCode.InternalServerError,
+                        subStatusCode: DataApiBuilderException.SubStatusCodes.UnexpectedError);
+                }
+
+                // Row didn't exist but INSERT returned no rows — create policy blocked it.
+                throw new DataApiBuilderException(
+                    message: DataApiBuilderException.AUTHORIZATION_FAILURE,
+                    statusCode: HttpStatusCode.Forbidden,
+                    subStatusCode: DataApiBuilderException.SubStatusCodes.DatabasePolicyFailure);
+            }
+
+            return dbResultSet;
+        }
+
         /// <summary>
         /// Determines if the saved default azure credential's access token is valid and not expired.
         /// </summary>

src/Core/Resolvers/PostgresQueryBuilder.cs

@@ -17,6 +17,7 @@ public class PostgresQueryBuilder : BaseSqlQueryBuilder, IQueryBuilder
         private const string UPSERT_IDENTIFIER_COLUMN_NAME = "___upsert_op___";
         private const string INSERT_UPSERT = "inserted";
         private const string UPDATE_UPSERT = "updated";
+        public const string COUNT_ROWS_WITH_GIVEN_PK = "cnt_rows_to_update";
 
         private static DbCommandBuilder _builder = new NpgsqlCommandBuilder();
 
@@ -67,11 +68,17 @@ public string Build(SqlQueryStructure structure)
         /// <inheritdoc />
         public string Build(SqlInsertStructure structure)
         {
-            string insertQuery = $"INSERT INTO {QuoteIdentifier(structure.DatabaseObject.SchemaName)}.{QuoteIdentifier(structure.DatabaseObject.Name)} ";
+            string tableName = $"{QuoteIdentifier(structure.DatabaseObject.SchemaName)}.{QuoteIdentifier(structure.DatabaseObject.Name)}";
+            string dbPolicyPredicates = JoinPredicateStrings(structure.GetDbPolicyForOperation(EntityActionOperation.Create));
+            string insertQuery = $"INSERT INTO {tableName} ";
+
             if (structure.InsertColumns.Any())
             {
-                insertQuery += $"({Build(structure.InsertColumns)}) " +
-                    $"VALUES ({string.Join(", ", (structure.Values))}) ";
+                string insertColumns = Build(structure.InsertColumns);
+                string insertValues = dbPolicyPredicates.Equals(BASE_PREDICATE)
+                    ? $"({insertColumns}) VALUES ({string.Join(", ", structure.Values)})"
+                    : $"({insertColumns}) SELECT {insertColumns} FROM (SELECT {string.Join(", ", structure.InsertColumns.Zip(structure.Values, (col, val) => $"{val} AS {QuoteIdentifier(col)}"))}) AS T WHERE {dbPolicyPredicates}";
+                insertQuery += insertValues;
             }
             else
             {
@@ -117,25 +124,50 @@ public string Build(SqlUpsertQueryStructure structure)
         {
             // https://stackoverflow.com/questions/42668720/check-if-postgres-query-inserted-or-updated-via-upsert
             // relying on xmax to detect insert vs update breaks for views
-            string updatePredicates = JoinPredicateStrings(Build(structure.Predicates), structure.GetDbPolicyForOperation(EntityActionOperation.Update));
-            string updateQuery = $"UPDATE {QuoteIdentifier(structure.DatabaseObject.SchemaName)}.{QuoteIdentifier(structure.DatabaseObject.Name)} " +
+            string tableName = $"{QuoteIdentifier(structure.DatabaseObject.SchemaName)}.{QuoteIdentifier(structure.DatabaseObject.Name)}";
+            string pkPredicates = Build(structure.Predicates);
+
+            // RS1: COUNT of rows matching PK (no policy) — used to distinguish
+            // "row doesn't exist" from "row exists but policy blocked" in the executor.
+            string countQuery = $"SELECT COUNT(*) AS {COUNT_ROWS_WITH_GIVEN_PK} FROM {tableName} WHERE {pkPredicates}";
+
+            string updatePredicates = JoinPredicateStrings(pkPredicates, structure.GetDbPolicyForOperation(EntityActionOperation.Update));
+            string updateQuery = $"UPDATE {tableName} " +
                 $"SET {Build(structure.UpdateOperations, ", ")} " +
                 $"WHERE {updatePredicates} " +
                 $"RETURNING {Build(structure.OutputColumns)}, '{UPDATE_UPSERT}' AS {UPSERT_IDENTIFIER_COLUMN_NAME}";
 
             if (structure.IsFallbackToUpdate)
             {
-                return updateQuery + ";";
+                // RS2: UPDATE only — no INSERT branch for autogen PK or missing required columns.
+                return $"{countQuery}; {updateQuery};";
             }
             else
             {
-                return $"WITH update_cte AS ( {updateQuery} ), insert_cte AS ( " +
-                    $"INSERT INTO {QuoteIdentifier(structure.DatabaseObject.SchemaName)}.{QuoteIdentifier(structure.DatabaseObject.Name)} ({Build(structure.InsertColumns)}) " +
-                    $"SELECT {string.Join(", ", (structure.Values))} " +
-                    $"WHERE NOT EXISTS (SELECT 1 FROM update_cte) " +
+                // INSERT only runs when row doesn't exist (pkPredicates match nothing)
+                // AND the create policy (if any) is satisfied.
+                string insertPredicates = JoinPredicateStrings(
+                    $"NOT EXISTS (SELECT 1 FROM {tableName} WHERE {pkPredicates})",
+                    structure.GetDbPolicyForOperation(EntityActionOperation.Create));
+
+                // Alias each value with its column name so that policy predicates referencing
+                // column names (e.g. "pieceid" != @param) can be resolved in the WHERE clause.
+                // Using SELECT ... FROM (SELECT @p1 AS col1, ...) AS T avoids both the VALUES(NULL)
+                // type inference issue and the unnamed-column resolution issue.
+                string namedValues = string.Join(", ",
+                    structure.InsertColumns.Zip(structure.Values,
+                        (col, val) => $"{val} AS {QuoteIdentifier(col)}"));
+
+                // RS2: CTE that attempts UPDATE first; falls through to INSERT only when row is absent.
+                string cteQuery = $"WITH update_cte AS ( {updateQuery} ), insert_cte AS ( " +
+                    $"INSERT INTO {tableName} ({Build(structure.InsertColumns)}) " +
+                    $"SELECT {Build(structure.InsertColumns)} FROM (SELECT {namedValues}) AS T " +
+                    $"WHERE {insertPredicates} " +
                     $"RETURNING {Build(structure.OutputColumns)}, '{INSERT_UPSERT}' AS {UPSERT_IDENTIFIER_COLUMN_NAME} ) " +
-                    $"SELECT {BuildListOfLabels(structure.OutputColumns)}, {UPSERT_IDENTIFIER_COLUMN_NAME} FROM update_cte UNION " +
+                    $"SELECT {BuildListOfLabels(structure.OutputColumns)}, {UPSERT_IDENTIFIER_COLUMN_NAME} FROM update_cte UNION ALL " +
                     $"SELECT {BuildListOfLabels(structure.OutputColumns)}, {UPSERT_IDENTIFIER_COLUMN_NAME} FROM insert_cte;";
+
+                return $"{countQuery}; {cteQuery}";
             }
         }
 

src/Core/Resolvers/SqlMutationEngine.cs

@@ -2015,6 +2015,7 @@ private async Task<DbResultSet?>
             IQueryBuilder queryBuilder = _queryManagerFactory.GetQueryBuilder(sqlMetadataProvider.GetDatabaseType());
             IQueryExecutor queryExecutor = _queryManagerFactory.GetQueryExecutor(sqlMetadataProvider.GetDatabaseType());
             string dataSourceName = _runtimeConfigProvider.GetConfig().GetDataSourceNameFromEntityName(entityName);
+            bool isFallbackToUpdate;
 
             if (operationType is EntityActionOperation.Upsert)
             {
@@ -2028,6 +2029,7 @@ private async Task<DbResultSet?>
                     incrementalUpdate: false);
                 queryString = queryBuilder.Build(upsertStructure);
                 queryParameters = upsertStructure.Parameters;
+                isFallbackToUpdate = upsertStructure.IsFallbackToUpdate;
             }
             else
             {
@@ -2041,6 +2043,7 @@ private async Task<DbResultSet?>
                     incrementalUpdate: true);
                 queryString = queryBuilder.Build(upsertIncrementalStructure);
                 queryParameters = upsertIncrementalStructure.Parameters;
+                isFallbackToUpdate = upsertIncrementalStructure.IsFallbackToUpdate;
             }
 
             string prettyPrintPk = "<" + string.Join(", ", context.PrimaryKeyValuePairs.Select(
@@ -2053,7 +2056,9 @@ private async Task<DbResultSet?>
                        queryExecutor.GetMultipleResultSetsIfAnyAsync,
                        dataSourceName,
                        GetHttpContext(),
-                       new List<string> { prettyPrintPk, entityName });
+                       sqlMetadataProvider.GetDatabaseType() is DatabaseType.PostgreSQL
+                           ? new List<string> { prettyPrintPk, entityName, isFallbackToUpdate.ToString() }
+                           : new List<string> { prettyPrintPk, entityName });
         }
 
         private Dictionary<string, object?> PrepareParameters(RestRequestContext context)
@@ -2264,7 +2269,7 @@ private void AuthorizeEntityAndFieldsForMutation(
         ///     }
         /// }
         /// </code>
-        /// 
+        ///
         /// Example 2 - Multiple items creation:
         /// <code>
         /// mutation {

src/Service.Tests/Snapshots/ConfigurationTests.TestReadingRuntimeConfigForPostgreSql.verified.txt

@@ -331,16 +331,19 @@
             Role: database_policy_tester,
             Actions: [
               {
-                Action: Update,
-                Policy: {
-                  Database: @item.pieceid ne 1
-                }
+                Action: Read
               },
               {
-                Action: Create
+                Action: Create,
+                Policy: {
+                  Database: @item.pieceid ne 6 and @item.piecesAvailable gt 0
+                }
               },
               {
-                Action: Read
+                Action: Update,
+                Policy: {
+                  Database: @item.pieceid ne 1
+                }
               }
             ]
           },

src/Service.Tests/SqlTests/RestApiTests/Patch/PostgreSqlPatchApiTests.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System;
 using System.Collections.Generic;
 using System.Threading.Tasks;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -166,6 +165,19 @@ SELECT to_jsonb(subq) AS data
                     ) AS subq
                 "
             },
+            {
+                "PatchOneInsertWithDatabasePolicy",
+                @"
+                    SELECT to_jsonb(subq) AS data
+                    FROM (
+                        SELECT categoryid, pieceid, ""categoryName"", ""piecesAvailable"", ""piecesRequired""
+                        FROM " + _Composite_NonAutoGenPK_TableName + @"
+                        WHERE categoryid = 0 AND pieceid = 7 AND ""categoryName"" = 'SciFi'
+                            AND ""piecesAvailable"" = 4 AND ""piecesRequired"" = 0
+                            AND (pieceid != 6 AND ""piecesAvailable"" > 0)
+                    ) AS subq
+                "
+            },
             {
                 "PatchOne_Update_Default_Test",
                 @"
@@ -335,27 +347,6 @@ await base.PatchOneViewBadRequestTest(
         }
 
         #region overridden tests
-
-        [TestMethod]
-        [Ignore]
-        public override Task PatchOneUpdateWithUnsatisfiedDatabasePolicy()
-        {
-            throw new NotImplementedException();
-        }
-
-        [TestMethod]
-        [Ignore]
-        public override Task PatchOneInsertWithUnsatisfiedDatabasePolicy()
-        {
-            throw new NotImplementedException();
-        }
-
-        [TestMethod]
-        [Ignore]
-        public override Task PatchOneInsertWithDatabasePolicy()
-        {
-            throw new NotImplementedException();
-        }
         #endregion
 
         #region Test Fixture Setup

src/Service.Tests/SqlTests/RestApiTests/Put/PostgreSqlPutApiTests.cs

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System;
 using System.Collections.Generic;
 using System.Threading.Tasks;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -85,6 +84,19 @@ SELECT to_jsonb(subq) AS data
                     ) AS subq
                 "
             },
+            {
+                "PutOneInsertWithDatabasePolicy",
+                @"
+                    SELECT to_jsonb(subq) AS data
+                    FROM (
+                        SELECT categoryid, pieceid, ""categoryName"", ""piecesAvailable"", ""piecesRequired""
+                        FROM " + _Composite_NonAutoGenPK_TableName + @"
+                        WHERE categoryid = 0 AND pieceid = 7 AND ""categoryName"" = 'SciFi'
+                            AND ""piecesAvailable"" = 4 AND ""piecesRequired"" = 0
+                            AND (pieceid != 6 AND ""piecesAvailable"" > 0)
+                    ) AS subq
+                "
+            },
             {
                 "PutOneUpdateAccessibleRowWithDatabasePolicy",
                 @"
@@ -434,27 +446,6 @@ public static async Task SetupAsync(TestContext context)
         #endregion
 
         #region overridden tests
-
-        [TestMethod]
-        [Ignore]
-        public override Task PutOneInsertWithDatabasePolicy()
-        {
-            throw new NotImplementedException();
-        }
-
-        [TestMethod]
-        [Ignore]
-        public override Task PutOneWithUnsatisfiedDatabasePolicy()
-        {
-            throw new NotImplementedException();
-        }
-
-        [TestMethod]
-        [Ignore]
-        public override Task PutOneInsertInTableWithFieldsInDbPolicyNotPresentInBody()
-        {
-            throw new NotImplementedException();
-        }
         #endregion
 
         [TestCleanup]