diff --git a/docs/v3/develop/settings-ref.mdx b/docs/v3/develop/settings-ref.mdx index fd5520ff4174..3e5fc2d6a2e1 100644 --- a/docs/v3/develop/settings-ref.mdx +++ b/docs/v3/develop/settings-ref.mdx @@ -871,10 +871,33 @@ Number of seconds a runner should wait between heartbeats for flow runs. **TOML dotted key path**: `runner.server` +--- +## SQLAlchemyConnectArgsSettings +Settings for controlling SQLAlchemy connection behavior; note that these settings only take effect when +using a PostgreSQL database. +### `application_name` +Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect backend. + +**Type**: `string | None` + +**Default**: `None` + +**TOML dotted key path**: `server.database.sqlalchemy.connect_args.application_name` + +**Supported environment variables**: +`PREFECT_SERVER_DATABASE_SQLALCHEMY_CONNECT_ARGS_APPLICATION_NAME` + --- ## SQLAlchemySettings Settings for controlling SQLAlchemy behavior; note that these settings only take effect when using a PostgreSQL database. +### `connect_args` +Settings for controlling SQLAlchemy connection behavior + +**Type**: [SQLAlchemyConnectArgsSettings](#sqlalchemyconnectargssettings) + +**TOML dotted key path**: `server.database.sqlalchemy.connect_args` + ### `pool_size` Controls connection pool size of database connection pools from the Prefect backend. @@ -1241,18 +1264,6 @@ A connection timeout, in seconds, applied to database connections. Defaults to ` **Supported environment variables**: `PREFECT_SERVER_DATABASE_CONNECTION_TIMEOUT`, `PREFECT_API_DATABASE_CONNECTION_TIMEOUT` -### `connection_app_name` -Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect backend. - -**Type**: `string | None` - -**Default**: `None` - -**TOML dotted key path**: `server.database.connection_app_name` - -**Supported environment variables**: -`PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME` - --- ## ServerDeploymentsSettings ### `concurrency_slot_wait_seconds` diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index e3f5db64419c..45ab9740729f 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -739,9 +739,37 @@ "title": "RunnerSettings", "type": "object" }, + "SQLAlchemyConnectArgsSettings": { + "description": "Settings for controlling SQLAlchemy connection behavior; note that these settings only take effect when\nusing a PostgreSQL database.", + "properties": { + "application_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect backend.", + "supported_environment_variables": [ + "PREFECT_SERVER_DATABASE_SQLALCHEMY_CONNECT_ARGS_APPLICATION_NAME" + ], + "title": "Application Name" + } + }, + "title": "SQLAlchemyConnectArgsSettings", + "type": "object" + }, "SQLAlchemySettings": { "description": "Settings for controlling SQLAlchemy behavior; note that these settings only take effect when\nusing a PostgreSQL database.", "properties": { + "connect_args": { + "$ref": "#/$defs/SQLAlchemyConnectArgsSettings", + "description": "Settings for controlling SQLAlchemy connection behavior", + "supported_environment_variables": [] + }, "pool_size": { "default": 5, "description": "Controls connection pool size of database connection pools from the Prefect backend.", @@ -1092,22 +1120,6 @@ "PREFECT_API_DATABASE_CONNECTION_TIMEOUT" ], "title": "Connection Timeout" - }, - "connection_app_name": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect backend.", - "supported_environment_variables": [ - "PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME" - ], - "title": "Connection App Name" } }, "title": "ServerDatabaseSettings", diff --git a/src/prefect/server/database/configurations.py b/src/prefect/server/database/configurations.py index b2055b92541b..5ff2fc8089fd 100644 --- a/src/prefect/server/database/configurations.py +++ b/src/prefect/server/database/configurations.py @@ -26,8 +26,6 @@ PREFECT_API_DATABASE_CONNECTION_TIMEOUT, PREFECT_API_DATABASE_ECHO, PREFECT_API_DATABASE_TIMEOUT, - PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME, - PREFECT_SQLALCHEMY_MAX_OVERFLOW, PREFECT_TESTING_UNIT_TEST_MODE, get_current_settings, ) @@ -135,10 +133,12 @@ def __init__( or get_current_settings().server.database.sqlalchemy.pool_size ) self.sqlalchemy_max_overflow: Optional[int] = ( - sqlalchemy_max_overflow or PREFECT_SQLALCHEMY_MAX_OVERFLOW.value() + sqlalchemy_max_overflow + or get_current_settings().server.database.sqlalchemy.max_overflow ) self.connection_app_name: Optional[str] = ( - connection_app_name or PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME.value() + connection_app_name + or get_current_settings().server.database.sqlalchemy.connect_args.application_name ) def unique_key(self) -> tuple[Hashable, ...]: @@ -209,9 +209,10 @@ async def engine(self) -> AsyncEngine: kwargs: dict[ str, Any ] = get_current_settings().server.database.sqlalchemy.model_dump( - mode="json" + mode="json", ) - connect_args: dict[str, Any] = dict() + connect_args: dict[str, Any] = kwargs.pop("connect_args") + app_name = connect_args.pop("application_name", None) if self.timeout is not None: connect_args["command_timeout"] = self.timeout @@ -219,9 +220,9 @@ async def engine(self) -> AsyncEngine: if self.connection_timeout is not None: connect_args["timeout"] = self.connection_timeout - if self.connection_app_name is not None: + if self.connection_app_name is not None or app_name is not None: connect_args["server_settings"] = dict( - application_name=self.connection_app_name + application_name=self.connection_app_name or app_name ) if connect_args: diff --git a/src/prefect/settings/models/server/database.py b/src/prefect/settings/models/server/database.py index b851a1db7b51..2c487386f6e5 100644 --- a/src/prefect/settings/models/server/database.py +++ b/src/prefect/settings/models/server/database.py @@ -15,6 +15,22 @@ from prefect.settings.base import PrefectBaseSettings, build_settings_config +class SQLAlchemyConnectArgsSettings(PrefectBaseSettings): + """ + Settings for controlling SQLAlchemy connection behavior; note that these settings only take effect when + using a PostgreSQL database. + """ + + model_config: ClassVar[SettingsConfigDict] = build_settings_config( + ("server", "database", "sqlalchemy", "connect_args") + ) + + application_name: Optional[str] = Field( + default=None, + description="Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect backend.", + ) + + class SQLAlchemySettings(PrefectBaseSettings): """ Settings for controlling SQLAlchemy behavior; note that these settings only take effect when @@ -25,6 +41,11 @@ class SQLAlchemySettings(PrefectBaseSettings): ("server", "database", "sqlalchemy") ) + connect_args: SQLAlchemyConnectArgsSettings = Field( + default_factory=SQLAlchemyConnectArgsSettings, + description="Settings for controlling SQLAlchemy connection behavior", + ) + pool_size: int = Field( default=5, description="Controls connection pool size of database connection pools from the Prefect backend.", @@ -194,11 +215,6 @@ class ServerDatabaseSettings(PrefectBaseSettings): ), ) - connection_app_name: Optional[str] = Field( - default=None, - description="Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect backend.", - ) - # handle deprecated fields def __getattribute__(self, name: str) -> Any: diff --git a/tests/test_settings.py b/tests/test_settings.py index 9e23cf67e8b8..deb2076a51f2 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -305,7 +305,6 @@ "test_value": timedelta(seconds=10), "legacy": True, }, - "PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME": {"test_value": "testingconn"}, "PREFECT_SERVER_DATABASE_CONNECTION_TIMEOUT": {"test_value": 10.0}, "PREFECT_SERVER_DATABASE_CONNECTION_URL": {"test_value": "sqlite:///"}, "PREFECT_SERVER_DATABASE_DRIVER": {"test_value": "sqlite+aiosqlite"}, @@ -315,6 +314,9 @@ "PREFECT_SERVER_DATABASE_NAME": {"test_value": "prefect"}, "PREFECT_SERVER_DATABASE_PASSWORD": {"test_value": "password"}, "PREFECT_SERVER_DATABASE_PORT": {"test_value": 5432}, + "PREFECT_SERVER_DATABASE_SQLALCHEMY_CONNECT_ARGS_APPLICATION_NAME": { + "test_value": "prefect" + }, "PREFECT_SERVER_DATABASE_SQLALCHEMY_MAX_OVERFLOW": {"test_value": 10}, "PREFECT_SERVER_DATABASE_SQLALCHEMY_POOL_RECYCLE": {"test_value": 10}, "PREFECT_SERVER_DATABASE_SQLALCHEMY_POOL_SIZE": {"test_value": 10},