|
2 | 2 | // Licensed under the MIT License. |
3 | 3 |
|
4 | 4 | using System.Data.Common; |
| 5 | +using System.Net; |
5 | 6 | using Azure.Core; |
6 | 7 | using Azure.DataApiBuilder.Config; |
7 | 8 | using Azure.DataApiBuilder.Config.ObjectModel; |
8 | 9 | using Azure.DataApiBuilder.Core.Configurations; |
9 | 10 | using Azure.DataApiBuilder.Core.Models; |
| 11 | +using Azure.DataApiBuilder.Service.Exceptions; |
10 | 12 | using Azure.Identity; |
11 | 13 | using Microsoft.AspNetCore.Http; |
12 | 14 | using Microsoft.Extensions.Logging; |
@@ -146,6 +148,87 @@ private static bool ShouldManagedIdentityAccessBeAttempted(NpgsqlConnectionStrin |
146 | 148 | return string.IsNullOrEmpty(builder.Password); |
147 | 149 | } |
148 | 150 |
|
| 151 | + /// <inheritdoc/> |
| 152 | + public override async Task<DbResultSet> GetMultipleResultSetsIfAnyAsync( |
| 153 | + DbDataReader dbDataReader, List<string>? args = null) |
| 154 | + { |
| 155 | + // RS1: COUNT of rows matching PK (no policy) — used to distinguish |
| 156 | + // "row doesn't exist" from "row exists but policy blocked". |
| 157 | + DbResultSet resultSetWithCountOfRowsWithGivenPk = await ExtractResultSetFromDbDataReaderAsync(dbDataReader); |
| 158 | + DbResultSetRow? resultSetRowWithCountOfRowsWithGivenPk = resultSetWithCountOfRowsWithGivenPk.Rows.FirstOrDefault(); |
| 159 | + int numOfRecordsWithGivenPK; |
| 160 | + bool isFallbackToUpdate; |
| 161 | + |
| 162 | + if (resultSetRowWithCountOfRowsWithGivenPk is not null && |
| 163 | + resultSetRowWithCountOfRowsWithGivenPk.Columns.TryGetValue(PostgresQueryBuilder.COUNT_ROWS_WITH_GIVEN_PK, out object? rowsWithGivenPK) && |
| 164 | + resultSetRowWithCountOfRowsWithGivenPk.Columns.TryGetValue(PostgresQueryBuilder.IS_FALLBACK_TO_UPDATE, out object? fallbackToUpdate)) |
| 165 | + { |
| 166 | + // PostgreSQL COUNT(*) returns Int64; convert to int. |
| 167 | + numOfRecordsWithGivenPK = Convert.ToInt32(rowsWithGivenPK!); |
| 168 | + isFallbackToUpdate = Convert.ToBoolean(fallbackToUpdate!); |
| 169 | + } |
| 170 | + else |
| 171 | + { |
| 172 | + throw new DataApiBuilderException( |
| 173 | + message: $"Neither insert nor update could be performed.", |
| 174 | + statusCode: HttpStatusCode.InternalServerError, |
| 175 | + subStatusCode: DataApiBuilderException.SubStatusCodes.UnexpectedError); |
| 176 | + } |
| 177 | + |
| 178 | + // RS2: UPDATE result, or UPDATE+INSERT CTE result. |
| 179 | + DbResultSet dbResultSet = await dbDataReader.NextResultAsync() |
| 180 | + ? await ExtractResultSetFromDbDataReaderAsync(dbDataReader) |
| 181 | + : throw new DataApiBuilderException( |
| 182 | + message: $"Neither insert nor update could be performed.", |
| 183 | + statusCode: HttpStatusCode.InternalServerError, |
| 184 | + subStatusCode: DataApiBuilderException.SubStatusCodes.UnexpectedError); |
| 185 | + |
| 186 | + if (numOfRecordsWithGivenPK == 1) // Row existed — we attempted an UPDATE. |
| 187 | + { |
| 188 | + if (dbResultSet.Rows.Count == 0) |
| 189 | + { |
| 190 | + // Row exists but UPDATE returned no rows — update policy blocked it. |
| 191 | + throw new DataApiBuilderException( |
| 192 | + message: DataApiBuilderException.AUTHORIZATION_FAILURE, |
| 193 | + statusCode: HttpStatusCode.Forbidden, |
| 194 | + subStatusCode: DataApiBuilderException.SubStatusCodes.DatabasePolicyFailure); |
| 195 | + } |
| 196 | + } |
| 197 | + else if (dbResultSet.Rows.Count == 0) |
| 198 | + { |
| 199 | + // If true, the row simply didn't exist — return 404 (same as MsSql's null-RS2 path). |
| 200 | + // If false, the INSERT ran but create policy blocked it — return 403. |
| 201 | + |
| 202 | + if (isFallbackToUpdate) |
| 203 | + { |
| 204 | + if (args is not null && args.Count > 1) |
| 205 | + { |
| 206 | + string prettyPrintPk = args[0]; |
| 207 | + string entityName = args[1]; |
| 208 | + |
| 209 | + throw new DataApiBuilderException( |
| 210 | + message: $"Cannot perform INSERT and could not find {entityName} " + |
| 211 | + $"with primary key {prettyPrintPk} to perform UPDATE on.", |
| 212 | + statusCode: HttpStatusCode.NotFound, |
| 213 | + subStatusCode: DataApiBuilderException.SubStatusCodes.ItemNotFound); |
| 214 | + } |
| 215 | + |
| 216 | + throw new DataApiBuilderException( |
| 217 | + message: $"Neither insert nor update could be performed.", |
| 218 | + statusCode: HttpStatusCode.InternalServerError, |
| 219 | + subStatusCode: DataApiBuilderException.SubStatusCodes.UnexpectedError); |
| 220 | + } |
| 221 | + |
| 222 | + // Row didn't exist but INSERT returned no rows — create policy blocked it. |
| 223 | + throw new DataApiBuilderException( |
| 224 | + message: DataApiBuilderException.AUTHORIZATION_FAILURE, |
| 225 | + statusCode: HttpStatusCode.Forbidden, |
| 226 | + subStatusCode: DataApiBuilderException.SubStatusCodes.DatabasePolicyFailure); |
| 227 | + } |
| 228 | + |
| 229 | + return dbResultSet; |
| 230 | + } |
| 231 | + |
149 | 232 | /// <summary> |
150 | 233 | /// Determines if the saved default azure credential's access token is valid and not expired. |
151 | 234 | /// </summary> |
|
0 commit comments