From 76b88d760a565b91cb8a6ea1abbeb3bc20636734 Mon Sep 17 00:00:00 2001
From: Daniel Voina <72201489+vs-dsva@users.noreply.github.com>
Date: Sun, 19 Oct 2025 12:32:26 +0300
Subject: [PATCH] Fix command timeout issues in DAB
---
 docs/postgresql-command-timeout.md            | 80 +++++++++++++++++++
 .../connection-string-timeout-example.json    | 56 +++++++++++++
 .../postgresql-command-timeout-example.json   | 59 ++++++++++++++
 schemas/dab.draft.schema.json                 | 15 +++-
 src/Config/ObjectModel/DataSource.cs          | 58 ++++++++++++++
 src/Config/RuntimeConfigLoader.cs             | 17 +++-
 .../Configuration/ConfigurationTests.cs       | 40 ++++++++++
 7 files changed, 321 insertions(+), 4 deletions(-)
 create mode 100644 docs/postgresql-command-timeout.md
 create mode 100644 samples/connection-string-timeout-example.json
 create mode 100644 samples/postgresql-command-timeout-example.json
diff --git a/docs/postgresql-command-timeout.md b/docs/postgresql-command-timeout.md
new file mode 100644
index 0000000000..c9c6467c5a
--- /dev/null
+++ b/docs/postgresql-command-timeout.md
@@ -0,0 +1,80 @@
+# PostgreSQL Command Timeout Configuration
+
+This document describes how to configure command timeout for PostgreSQL data sources in Data API Builder.
+
+## Overview
+
+Data API Builder now supports configuring PostgreSQL command timeout through the `command-timeout` option in the data source configuration. This feature allows you to override the default command timeout for all PostgreSQL queries executed by Data API Builder.
+
+## Configuration
+
+Add the `command-timeout` option to your PostgreSQL data source configuration:
+
+```json
+{
+  "data-source": {
+    "database-type": "postgresql",
+    "connection-string": "Host=localhost;Database=mydb;Username=user;Password=pass;",
+    "options": {
+      "command-timeout": 60
+    }
+  }
+}
+```
+
+### Parameters
+
+- **command-timeout**: Integer value representing the timeout in seconds
+  - **Type**: `integer`
+  - **Minimum**: `0`
+  - **Default**: `30` (if not specified)
+  - **Description**: Sets the wait time (in seconds) before terminating the attempt to execute a command and generating an error
+
+### Behavior
+
+1. **Override**: The `command-timeout` value from the configuration will override any `CommandTimeout` parameter present in the connection string
+2. **Precedence**: Configuration file setting takes priority over connection string setting
+3. **Scope**: Applies to all PostgreSQL queries executed through Data API Builder
+
+### Example
+
+```json
+{
+  "$schema": "schemas/dab.draft.schema.json",
+  "data-source": {
+    "database-type": "postgresql",
+    "connection-string": "Host=localhost;Database=bookstore;Username=postgres;Password=password;",
+    "options": {
+      "command-timeout": 120
+    }
+  },
+  "runtime": {
+    "rest": { "enabled": true, "path": "/api" },
+    "graphql": { "enabled": true, "path": "/graphql" }
+  },
+  "entities": {
+    "Book": {
+      "source": { "object": "books", "type": "table" },
+      "permissions": [
+        { "role": "anonymous", "actions": [{ "action": "*" }] }
+      ]
+    }
+  }
+}
+```
+
+In this example, all PostgreSQL queries will have a 120-second timeout, regardless of any `CommandTimeout` value in the connection string.
+
+## Implementation Details
+
+The feature is implemented by:
+
+1. **Schema Validation**: The JSON schema validates the `command-timeout` parameter
+2. **Options Parsing**: The `PostgreSqlOptions` class parses the timeout value from various data types (integer, string, JsonElement)
+3. **Connection String Processing**: The timeout is applied to the Npgsql connection string builder during connection string normalization
+4. **Override Logic**: Configuration values take precedence over existing connection string parameters
+
+## Related
+
+- See `samples/postgresql-command-timeout-example.json` for a complete working example
+- For other database types, command timeout can be configured directly in the connection string
\ No newline at end of file
diff --git a/samples/connection-string-timeout-example.json b/samples/connection-string-timeout-example.json
new file mode 100644
index 0000000000..7d28f229b9
--- /dev/null
+++ b/samples/connection-string-timeout-example.json
@@ -0,0 +1,56 @@
+{
+  "$schema": "../schemas/dab.draft.schema.json",
+  "data-source": {
+    "database-type": "postgresql",
+    "connection-string": "Host=localhost;Port=5432;Username=postgres;Password=password;Database=bookshelf;CommandTimeout=120"
+  },
+  "runtime": {
+    "rest": {
+      "enabled": true,
+      "path": "/api"
+    },
+    "graphql": {
+      "enabled": true,
+      "path": "/graphql",
+      "allow-introspection": true
+    },
+    "host": {
+      "cors": {
+        "origins": ["*"],
+        "allow-credentials": false
+      },
+      "authentication": {
+        "provider": "StaticWebApps"
+      },
+      "mode": "development"
+    }
+  },
+  "entities": {
+    "Book": {
+      "source": {
+        "object": "books",
+        "type": "table"
+      },
+      "graphql": {
+        "enabled": true,
+        "type": {
+          "singular": "Book",
+          "plural": "Books"
+        }
+      },
+      "rest": {
+        "enabled": true
+      },
+      "permissions": [
+        {
+          "role": "anonymous",
+          "actions": [
+            {
+              "action": "*"
+            }
+          ]
+        }
+      ]
+    }
+  }
+}
\ No newline at end of file
diff --git a/samples/postgresql-command-timeout-example.json b/samples/postgresql-command-timeout-example.json
new file mode 100644
index 0000000000..d64e3d642c
--- /dev/null
+++ b/samples/postgresql-command-timeout-example.json
@@ -0,0 +1,59 @@
+{
+  "$schema": "../schemas/dab.draft.schema.json",
+  "data-source": {
+    "database-type": "postgresql",
+    "connection-string": "Host=localhost;Database=bookstore;Username=postgres;Password=password;",
+    "options": {
+      "command-timeout": 60
+    }
+  },
+  "runtime": {
+    "rest": {
+      "enabled": true,
+      "path": "/api"
+    },
+    "graphql": {
+      "enabled": true,
+      "path": "/graphql",
+      "allow-introspection": true
+    },
+    "host": {
+      "cors": {
+        "origins": ["*"],
+        "allow-credentials": false
+      },
+      "authentication": {
+        "provider": "StaticWebApps"
+      },
+      "mode": "development"
+    }
+  },
+  "entities": {
+    "Book": {
+      "source": {
+        "object": "books",
+        "type": "table"
+      },
+      "graphql": {
+        "enabled": true,
+        "type": {
+          "singular": "Book",
+          "plural": "Books"
+        }
+      },
+      "rest": {
+        "enabled": true
+      },
+      "permissions": [
+        {
+          "role": "anonymous",
+          "actions": [
+            {
+              "action": "*"
+            }
+          ]
+        }
+      ]
+    }
+  }
+}
\ No newline at end of file
diff --git a/schemas/dab.draft.schema.json b/schemas/dab.draft.schema.json
index b348ac4a4f..63691daa34 100644
--- a/schemas/dab.draft.schema.json
+++ b/schemas/dab.draft.schema.json
@@ -115,6 +115,12 @@
                   "set-session-context": {
                     "type": "boolean",
                     "description": "Enable sending data to MsSql using session context"
+                  },
+                  "command-timeout": {
+                    "type": "integer",
+                    "description": "Command timeout in seconds for database operations",
+                    "minimum": 0,
+                    "default": 30
                   }
                 }
               }
@@ -135,7 +141,14 @@
               "options": {
                 "type": "object",
                 "additionalProperties": false,
-                "properties": {},
+                "properties": {
+                  "command-timeout": {
+                    "type": "integer",
+                    "description": "Command timeout in seconds for database operations",
+                    "minimum": 0,
+                    "default": 30
+                  }
+                },
                 "required": []
               }
             }
diff --git a/src/Config/ObjectModel/DataSource.cs b/src/Config/ObjectModel/DataSource.cs
index d1a2456ef9..0cf9f303aa 100644
--- a/src/Config/ObjectModel/DataSource.cs
+++ b/src/Config/ObjectModel/DataSource.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Text.Json;
 using System.Text.Json.Serialization;
 using Azure.DataApiBuilder.Config.HealthCheck;
 using Azure.DataApiBuilder.Config.NamingPolicies;
@@ -69,6 +70,12 @@ public int DatasourceThresholdMs
                 SetSessionContext: ReadBoolOption(namingPolicy.ConvertName(nameof(MsSqlOptions.SetSessionContext))));
         }
 
+        if (typeof(TOptionType).IsAssignableFrom(typeof(PostgreSqlOptions)))
+        {
+            return (TOptionType)(object)new PostgreSqlOptions(
+                CommandTimeout: ReadIntOption(namingPolicy.ConvertName("command-timeout")));
+        }
+
         throw new NotSupportedException($"The type {typeof(TOptionType).FullName} is not a supported strongly typed options object");
     }
 
@@ -92,6 +99,52 @@ private bool ReadBoolOption(string option)
         return false;
     }
 
+    private int? ReadIntOption(string option)
+    {
+        if (Options is not null && Options.TryGetValue(option, out object? value))
+        {
+            if (value is int intValue)
+            {
+                return intValue;
+            }
+            else if (value is string stringValue && int.TryParse(stringValue, out int parsedValue))
+            {
+                return parsedValue;
+            }
+            else if (value is JsonElement jsonElement && jsonElement.ValueKind == JsonValueKind.Number && jsonElement.TryGetInt32(out int jsonValue))
+            {
+                return jsonValue;
+            }
+        }
+
+        return null;
+    }
+
+    /// 
+    /// Gets the command timeout value from the options.
+    /// 
+    /// The command timeout in seconds, or 30 (default) if not specified.
+    public int GetCommandTimeout()
+    {
+        if (Options is not null && Options.TryGetValue("command-timeout", out object? value))
+        {
+            if (value is int intValue)
+            {
+                return intValue;
+            }
+            else if (value is long longValue && longValue <= int.MaxValue && longValue >= int.MinValue)
+            {
+                return (int)longValue;
+            }
+            else if (value is string stringValue && int.TryParse(stringValue, out int parsedValue))
+            {
+                return parsedValue;
+            }
+        }
+
+        return 30; // default command timeout
+    }
+
     [JsonIgnore]
     public string DatabaseTypeNotSupportedMessage => $"The provided database-type value: {DatabaseType} is currently not supported. Please check the configuration file.";
 }
@@ -111,3 +164,8 @@ public record CosmosDbNoSQLDataSourceOptions(string? Database, string? Container
 /// Options for MsSql database.
 /// 
 public record MsSqlOptions(bool SetSessionContext = true) : IDataSourceOptions;
+
+/// 
+/// Options for PostgreSQL database.
+/// 
+public record PostgreSqlOptions(int? CommandTimeout = null) : IDataSourceOptions;
diff --git a/src/Config/RuntimeConfigLoader.cs b/src/Config/RuntimeConfigLoader.cs
index f78c32ebc1..1935d77262 100644
--- a/src/Config/RuntimeConfigLoader.cs
+++ b/src/Config/RuntimeConfigLoader.cs
@@ -186,7 +186,7 @@ public static bool TryParseConfig(string json,
                 }
                 else if (ds.DatabaseType is DatabaseType.PostgreSQL && replaceEnvVar)
                 {
-                    updatedConnection = GetPgSqlConnectionStringWithApplicationName(connectionValue);
+                    updatedConnection = GetPgSqlConnectionStringWithApplicationName(connectionValue, ds);
                 }
 
                 ds = ds with { ConnectionString = updatedConnection };
@@ -331,8 +331,9 @@ internal static string GetConnectionStringWithApplicationName(string connectionS
     /// else add the Application Name property with DataApiBuilder Application Name based on hosted/oss platform.
     /// 
     /// Connection string for connecting to database.
-    /// Updated connection string with `Application Name` property.
-    internal static string GetPgSqlConnectionStringWithApplicationName(string connectionString)
+    /// The data source configuration, used to get command timeout override.
+    /// Updated connection string with `Application Name` property and command timeout.
+    internal static string GetPgSqlConnectionStringWithApplicationName(string connectionString, DataSource? dataSource = null)
     {
         // If the connection string is null, empty, or whitespace, return it as is.
         if (string.IsNullOrWhiteSpace(connectionString))
@@ -369,6 +370,16 @@ internal static string GetPgSqlConnectionStringWithApplicationName(string connec
             connectionStringBuilder.ApplicationName += $",{applicationName}";
         }
 
+        // Apply command timeout from data source configuration if specified (overrides connection string value)
+        if (dataSource?.Options is not null)
+        {
+            PostgreSqlOptions? pgOptions = dataSource.GetTypedOptions();
+            if (pgOptions?.CommandTimeout is not null)
+            {
+                connectionStringBuilder.CommandTimeout = pgOptions.CommandTimeout.Value;
+            }
+        }
+
         // Return the updated connection string.
         return connectionStringBuilder.ConnectionString;
     }
diff --git a/src/Service.Tests/Configuration/ConfigurationTests.cs b/src/Service.Tests/Configuration/ConfigurationTests.cs
index 0be24fa886..7baea82872 100644
--- a/src/Service.Tests/Configuration/ConfigurationTests.cs
+++ b/src/Service.Tests/Configuration/ConfigurationTests.cs
@@ -5462,6 +5462,46 @@ public static EntityPermission GetMinimalPermissionConfig(string roleName)
             );
         }
 
+        /// 
+        /// Test that command timeout is properly read from data source options and applied to connection string.
+        /// 
+        [TestMethod]
+        public void TestCommandTimeoutFromDataSourceOptions()
+        {
+            // Test SQL Server
+            Dictionary mssqlOptions = new()
+            {
+                { "command-timeout", 60 },
+                { "set-session-context", true }
+            };
+            DataSource mssqlDataSource = new(DatabaseType.MSSQL, "Server=localhost;Database=test;", mssqlOptions);
+            Assert.AreEqual(60, mssqlDataSource.GetCommandTimeout());
+
+            // Test PostgreSQL
+            Dictionary pgOptions = new()
+            {
+                { "command-timeout", 120 }
+            };
+            DataSource pgDataSource = new(DatabaseType.PostgreSQL, "Host=localhost;Database=test;", pgOptions);
+            Assert.AreEqual(120, pgDataSource.GetCommandTimeout());
+
+            // Test MySQL
+            Dictionary mysqlOptions = new()
+            {
+                { "command-timeout", 45 }
+            };
+            DataSource mysqlDataSource = new(DatabaseType.MySQL, "Server=localhost;Database=test;", mysqlOptions);
+            Assert.AreEqual(45, mysqlDataSource.GetCommandTimeout());
+
+            // Test default value when not specified
+            DataSource defaultDataSource = new(DatabaseType.MSSQL, "Server=localhost;Database=test;", new());
+            Assert.AreEqual(30, defaultDataSource.GetCommandTimeout());
+
+            // Test null options
+            DataSource nullOptionsDataSource = new(DatabaseType.MSSQL, "Server=localhost;Database=test;", null);
+            Assert.AreEqual(30, nullOptionsDataSource.GetCommandTimeout());
+        }
+
         /// 
         /// Reads configuration file for defined environment to acquire the connection string.
         /// CI/CD Pipelines and local environments may not have connection string set as environment variable.