GH-46315 [C#] Apache Arrow Flight Middleware

csharp/src/Apache.Arrow.Flight/Middleware/ClientCookieMiddleware.cs

@@ -0,0 +1,76 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System.Collections.Generic;
+using Apache.Arrow.Flight.Middleware.Interfaces;
+using Apache.Arrow.Flight.Middleware.Models;
+using Microsoft.Extensions.Logging;
+
+namespace Apache.Arrow.Flight.Middleware;
+
+public class ClientCookieMiddleware : IFlightClientMiddleware
+{
+    private readonly ClientCookieMiddlewareFactory _factory;
+    private readonly ILogger<ClientCookieMiddleware> _logger;
+    private const string SET_COOKIE_HEADER = "Set-cookie";
+    private const string COOKIE_HEADER = "Cookie";
+
+    public ClientCookieMiddleware(ClientCookieMiddlewareFactory factory,
+        ILogger<ClientCookieMiddleware> logger)
+    {
+        _factory = factory;
+        _logger = logger;
+    }
+
+    public void OnBeforeSendingHeaders(ICallHeaders outgoingHeaders)
+    {
+        var cookieValue = GetValidCookiesAsString();
+        if (!string.IsNullOrEmpty(cookieValue))
+        {
+            outgoingHeaders.Insert(COOKIE_HEADER, cookieValue);
+        }
+
+        _logger.LogInformation("Sending Headers: " + string.Join(", ", outgoingHeaders.Keys));
+    }
+
+    public void OnHeadersReceived(ICallHeaders incomingHeaders)
+    {
+        var setCookieHeaders = incomingHeaders.GetAll(SET_COOKIE_HEADER);
+        _factory.UpdateCookies(setCookieHeaders);
+        _logger.LogInformation("Received Headers: " + string.Join(", ", incomingHeaders.Keys));
+    }
+
+    public void OnCallCompleted(CallStatus status)
+    {
+        _logger.LogInformation($"Call completed with: {status.Code} ({status.Description})");
+    }
+
+    private string GetValidCookiesAsString()
+    {
+        var cookieList = new List<string>();
+        foreach (var entry in _factory.Cookies)
+        {
+            if (entry.Value.Expired)
+            {
+                _factory.Cookies.TryRemove(entry.Key, out _);
+            }
+            else
+            {
+                cookieList.Add(entry.Value.ToString());
+            }
+        }
+        return string.Join("; ", cookieList);
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/ClientCookieMiddlewareFactory.cs

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Net;
+using Apache.Arrow.Flight.Middleware.Extensions;
+using Apache.Arrow.Flight.Middleware.Interfaces;
+using Microsoft.Extensions.Logging;
+using CallInfo = Apache.Arrow.Flight.Middleware.Models.CallInfo;
+
+namespace Apache.Arrow.Flight.Middleware;
+
+public class ClientCookieMiddlewareFactory : IFlightClientMiddlewareFactory
+{
+    public readonly ConcurrentDictionary<string, Cookie> Cookies = new(StringComparer.OrdinalIgnoreCase);
+    private readonly ILoggerFactory _loggerFactory;
+
+    public ClientCookieMiddlewareFactory(ILoggerFactory loggerFactory)
+    {
+        _loggerFactory = loggerFactory;
+    }
+
+    public IFlightClientMiddleware OnCallStarted(CallInfo callInfo)
+    {
+        var logger = _loggerFactory.CreateLogger<ClientCookieMiddleware>();
+        return new ClientCookieMiddleware(this, logger);
+    }
+
+    internal void UpdateCookies(IEnumerable<string> newCookieHeaderValues)
+    {
+        foreach (var headerValue in newCookieHeaderValues)
+        {
+            try
+            {
+                var parsedCookies = headerValue.ParseHeader();
+                foreach (var parsedCookie in parsedCookies)
+                {
+                    var nameLc = parsedCookie.Name.ToLower(CultureInfo.InvariantCulture);
+                    if (parsedCookie.Expired)
+                    {
+                        Cookies.TryRemove(nameLc, out _);
+                    }
+                    else
+                    {
+                        Cookies[nameLc] = parsedCookie;
+                    }
+                }
+            }
+            catch (FormatException ex)
+            {
+                var logger = _loggerFactory.CreateLogger<ClientCookieMiddleware>();
+                logger.LogWarning(ex, "Skipping malformed Set-Cookie header: '{HeaderValue}'", headerValue);
+            }
+        }
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/Extensions/CookieExtensions.cs

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net;
+
+namespace Apache.Arrow.Flight.Middleware.Extensions;
+
+public static class CookieExtensions
+{
+    public static IEnumerable<Cookie> ParseHeader(this string headers)
+    {
+        if (string.IsNullOrEmpty(headers))
+            return System.Array.Empty<Cookie>();
+
+        var cookies = new List<Cookie>();
+        var segments = headers.Split([';'], StringSplitOptions.RemoveEmptyEntries);
+
+        if (segments.Length == 0) return cookies;
+
+        var nameValue = segments[0].Split(['='], 2);
+        if (nameValue.Length == 2)
+        {
+            var cookie = new Cookie(nameValue[0].Trim(), nameValue[1].Trim());
+
+            foreach (var segment in segments.Skip(1))
+            {
+                var trimmedSegment = segment.Trim();
+                if (trimmedSegment.StartsWith("Expires=", StringComparison.OrdinalIgnoreCase))
+                {
+                    var value = trimmedSegment.Substring("Expires=".Length).Trim();
+
+                    if (!DateTimeOffset.TryParseExact(
+                            value,
+                            "R",
+                            System.Globalization.CultureInfo.InvariantCulture,
+                            System.Globalization.DateTimeStyles.None,
+                            out var expires))
+                    {
+                        if (DateTimeOffset.TryParse(value, out expires))
+                        {
+                            cookie.Expires = expires.UtcDateTime;
+                        }
+                    }
+                    else
+                    {
+                        cookie.Expires = expires.UtcDateTime;
+                    }
+                }
+            }
+
+            cookies.Add(cookie);
+        }
+
+        return cookies;
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/Extensions/FlightMethodParser.cs

@@ -0,0 +1,63 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Apache.Arrow.Flight.Middleware.Models;
+
+namespace Apache.Arrow.Flight.Middleware.Extensions;
+
+public static class FlightMethodParser
+{
+    /// <summary>
+    /// Parses the gRPC full method name (e.g., "/arrow.flight.protocol.FlightService/DoGet")
+    /// and maps it to a known FlightMethod.
+    /// </summary>
+    /// <param name="fullMethodName">gRPC method name</param>
+    /// <returns>Parsed FlightMethod</returns>
+    public static FlightMethod Parse(string fullMethodName)
+    {
+        if (string.IsNullOrWhiteSpace(fullMethodName))
+            return FlightMethod.Unknown;
+
+        var parts = fullMethodName.Split('/');
+        if (parts.Length < 2)
+            return FlightMethod.Unknown;
+
+        var methodName = parts[parts.Length - 1];
+
+        return methodName switch
+        {
+            "Handshake" => FlightMethod.Handshake,
+            "ListFlights" => FlightMethod.ListFlights,
+            "GetFlightInfo" => FlightMethod.GetFlightInfo,
+            "GetSchema" => FlightMethod.GetSchema,
+            "DoGet" => FlightMethod.DoGet,
+            "DoPut" => FlightMethod.DoPut,
+            "DoExchange" => FlightMethod.DoExchange,
+            "DoAction" => FlightMethod.DoAction,
+            "ListActions" => FlightMethod.ListActions,
+            "CancelFlightInfo" => FlightMethod.CancelFlightInfo,
+            _ => FlightMethod.Unknown
+        };
+    }
+
+    public static string ParseMethodName(string fullMethodName)
+    {
+        if (string.IsNullOrWhiteSpace(fullMethodName))
+            return "Unknown";
+
+        var parts = fullMethodName.Split('/');
+        return parts.Length >= 2 ? parts[parts.Length - 1] : "Unknown";
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/Extensions/StatusUtils.cs

@@ -0,0 +1,58 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Apache.Arrow.Flight.Middleware.Models;
+using Grpc.Core;
+
+namespace Apache.Arrow.Flight.Middleware.Extensions;
+
+public static class StatusUtils
+{
+    public static CallStatus FromGrpcStatusAndTrailers(Status status, Metadata trailers)
+    {
+        var code = FromGrpcStatusCode(status.StatusCode);
+        return new CallStatus(
+            code,
+            status.StatusCode != StatusCode.OK ? new RpcException(status, trailers) : null,
+            status.Detail,
+            trailers
+        );
+    }
+
+    public static FlightStatusCode FromGrpcStatusCode(StatusCode grpcCode)
+    {
+        return grpcCode switch
+        {
+            StatusCode.OK => FlightStatusCode.Ok,
+            StatusCode.Cancelled => FlightStatusCode.Cancelled,
+            StatusCode.Unknown => FlightStatusCode.Unknown,
+            StatusCode.InvalidArgument => FlightStatusCode.InvalidArgument,
+            StatusCode.DeadlineExceeded => FlightStatusCode.DeadlineExceeded,
+            StatusCode.NotFound => FlightStatusCode.NotFound,
+            StatusCode.AlreadyExists => FlightStatusCode.AlreadyExists,
+            StatusCode.PermissionDenied => FlightStatusCode.PermissionDenied,
+            StatusCode.Unauthenticated => FlightStatusCode.Unauthenticated,
+            StatusCode.ResourceExhausted => FlightStatusCode.ResourceExhausted,
+            StatusCode.FailedPrecondition => FlightStatusCode.FailedPrecondition,
+            StatusCode.Aborted => FlightStatusCode.Aborted,
+            StatusCode.OutOfRange => FlightStatusCode.OutOfRange,
+            StatusCode.Unimplemented => FlightStatusCode.Unimplemented,
+            StatusCode.Internal => FlightStatusCode.Internal,
+            StatusCode.Unavailable => FlightStatusCode.Unavailable,
+            StatusCode.DataLoss => FlightStatusCode.DataLoss,
+            _ => FlightStatusCode.Unknown
+        };
+    }
+}