Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds new connection application name setting #16690

Merged
merged 5 commits into from
Jan 13, 2025
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion docs/v3/develop/settings-ref.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -1183,7 +1183,7 @@ A connection timeout, in seconds, applied to database connections. Defaults to `
`PREFECT_SERVER_DATABASE_CONNECTION_TIMEOUT`, `PREFECT_API_DATABASE_CONNECTION_TIMEOUT`

### `sqlalchemy_pool_size`
Controls connection pool size when using a PostgreSQL database with the Prefect API. If not set, the default SQLAlchemy pool size will be used.
Controls connection pool size of database connection pools from the Prefect API. If not set, the default SQLAlchemy pool size will be used.

**Type**: `integer | None`

Expand All @@ -1194,6 +1194,18 @@ Controls connection pool size when using a PostgreSQL database with the Prefect
**Supported environment variables**:
`PREFECT_SERVER_DATABASE_SQLALCHEMY_POOL_SIZE`, `PREFECT_SQLALCHEMY_POOL_SIZE`

### `connection_app_name`
Controls the application_name field for connections opened from the connection pool when using a PostgreSQL database with the Prefect API.

**Type**: `string | None`

**Default**: `None`

**TOML dotted key path**: `server.database.connection_app_name`

**Supported environment variables**:
`PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME`

### `sqlalchemy_max_overflow`
Controls maximum overflow of the connection pool when using a PostgreSQL database with the Prefect API. If not set, the default SQLAlchemy maximum overflow value will be used.

Expand Down
18 changes: 17 additions & 1 deletion schemas/settings.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -1046,13 +1046,29 @@
}
],
"default": null,
"description": "Controls connection pool size when using a PostgreSQL database with the Prefect API. If not set, the default SQLAlchemy pool size will be used.",
"description": "Controls connection pool size of database connection pools from the Prefect API. If not set, the default SQLAlchemy pool size will be used.",
"supported_environment_variables": [
"PREFECT_SERVER_DATABASE_SQLALCHEMY_POOL_SIZE",
"PREFECT_SQLALCHEMY_POOL_SIZE"
],
"title": "Sqlalchemy Pool Size"
},
"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 API.",
"supported_environment_variables": [
"PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME"
],
"title": "Connection App Name"
},
"sqlalchemy_max_overflow": {
"anyOf": [
{
Expand Down
16 changes: 16 additions & 0 deletions src/prefect/server/database/configurations.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
PREFECT_API_DATABASE_CONNECTION_TIMEOUT,
PREFECT_API_DATABASE_ECHO,
PREFECT_API_DATABASE_TIMEOUT,
PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME,
PREFECT_SQLALCHEMY_MAX_OVERFLOW,
PREFECT_SQLALCHEMY_POOL_SIZE,
PREFECT_TESTING_UNIT_TEST_MODE,
Expand Down Expand Up @@ -121,6 +122,7 @@ def __init__(
connection_timeout: Optional[float] = None,
sqlalchemy_pool_size: Optional[int] = None,
sqlalchemy_max_overflow: Optional[int] = None,
connection_app_name: Optional[str] = None,
) -> None:
self.connection_url = connection_url
self.echo: bool = echo or PREFECT_API_DATABASE_ECHO.value()
Expand All @@ -134,6 +136,9 @@ def __init__(
self.sqlalchemy_max_overflow: Optional[int] = (
sqlalchemy_max_overflow or PREFECT_SQLALCHEMY_MAX_OVERFLOW.value()
)
self.connection_app_name: Optional[str] = (
connection_app_name or PREFECT_SERVER_DATABASE_CONNECTION_APP_NAME.value()
)

def unique_key(self) -> tuple[Hashable, ...]:
"""
Expand Down Expand Up @@ -210,6 +215,11 @@ 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:
connect_args["server_settings"] = dict(
application_name=self.connection_app_name
)

if connect_args:
kwargs["connect_args"] = connect_args

Expand Down Expand Up @@ -337,6 +347,12 @@ async def engine(self) -> AsyncEngine:
if self.timeout is not None:
kwargs["connect_args"] = dict(timeout=self.timeout)

if self.sqlalchemy_pool_size is not None:
kwargs["pool_size"] = self.sqlalchemy_pool_size

if self.sqlalchemy_max_overflow is not None:
kwargs["max_overflow"] = self.sqlalchemy_max_overflow

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, I added these kwargs into our connection pools for non-memory SQLite as well; I couldn't find any reason why we should avoid configuring the pool for a SQLite DB.

# use `named` paramstyle for sqlite instead of `qmark` in very rare
# circumstances, we've seen aiosqlite pass parameters in the wrong
# order; by using named parameters we avoid this issue
Expand Down
11 changes: 10 additions & 1 deletion src/prefect/settings/models/server/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,14 +148,23 @@ class ServerDatabaseSettings(PrefectBaseSettings):

sqlalchemy_pool_size: Optional[int] = Field(
default=None,
description="Controls connection pool size when using a PostgreSQL database with the Prefect API. If not set, the default SQLAlchemy pool size will be used.",
description="Controls connection pool size of database connection pools from the Prefect API. If not set, the default SQLAlchemy pool size will be used.",
validation_alias=AliasChoices(
AliasPath("sqlalchemy_pool_size"),
"prefect_server_database_sqlalchemy_pool_size",
"prefect_sqlalchemy_pool_size",
),
)

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 API.",
validation_alias=AliasChoices(
AliasPath("connection_app_name"),
"prefect_server_database_connection_app_name",
),
cicdw marked this conversation as resolved.
Show resolved Hide resolved
)

sqlalchemy_max_overflow: Optional[int] = Field(
default=None,
description="Controls maximum overflow of the connection pool when using a PostgreSQL database with the Prefect API. If not set, the default SQLAlchemy maximum overflow value will be used.",
Expand Down
1 change: 1 addition & 0 deletions tests/test_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,7 @@
"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"},
Expand Down
Loading