fix(storage): use dynamic schema_getter in PostgreSQLFileStorage for multi-tenant (#440)

PostgreSQLFileStorage was initialized once at startup with a static
schema value. Since get_current_schema() returns the default schema at
init time, multi-tenant requests always queried the wrong schema,
causing "relation file_storage does not exist" errors.

Replace static schema with schema_getter callable (same pattern used
by BrokerTaskBackend since #208) so the schema is resolved dynamically
per-request via contextvars.
This commit is contained in:
And#ocean 2026-02-25 17:19:08 +03:00 committed by GitHub
parent 4b328a9cb3
commit 86d8ac08b1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 22 additions and 6 deletions

View file

@ -1366,7 +1366,7 @@ class MemoryEngine(MemoryEngineInterface):
self._file_storage = create_file_storage(
storage_type=config.file_storage_type,
pool_getter=lambda: self._pool,
schema=get_current_schema() if get_current_schema() != config.database_schema else None,
schema_getter=get_current_schema,
)
logger.debug(f"File storage initialized ({config.file_storage_type})")

View file

@ -12,6 +12,7 @@ def create_file_storage(
storage_type: str,
pool_getter: Callable | None = None,
schema: str | None = None,
schema_getter: Callable | None = None,
**kwargs,
) -> FileStorage:
"""
@ -20,7 +21,8 @@ def create_file_storage(
Args:
storage_type: "native" (PostgreSQL BYTEA) or "s3" (S3-compatible object storage)
pool_getter: Database pool getter (required for native)
schema: Database schema (for native multi-tenant)
schema: Static database schema (for native single-tenant)
schema_getter: Callable returning current schema at query time (for native multi-tenant)
**kwargs: Additional args passed to storage backend
Returns:
@ -32,7 +34,7 @@ def create_file_storage(
if storage_type == "native":
if not pool_getter:
raise ValueError("pool_getter required for native (PostgreSQL) storage")
return PostgreSQLFileStorage(pool_getter=pool_getter, schema=schema)
return PostgreSQLFileStorage(pool_getter=pool_getter, schema=schema, schema_getter=schema_getter)
elif storage_type == "s3":
from ...config import get_config
from .s3 import S3FileStorage

View file

@ -40,16 +40,30 @@ class PostgreSQLFileStorage(FileStorage):
For production/scale, consider S3FileStorage instead.
"""
def __init__(self, pool_getter: Callable[[], "asyncpg.Pool"], schema: str | None = None):
def __init__(
self,
pool_getter: Callable[[], "asyncpg.Pool"],
schema: str | None = None,
schema_getter: Callable[[], str] | None = None,
):
"""
Initialize PostgreSQL file storage.
Args:
pool_getter: Function that returns asyncpg connection pool
schema: Database schema (for multi-tenant support)
schema: Static database schema (fallback for single-tenant / tests)
schema_getter: Callable returning current schema at query time (for multi-tenant)
"""
self._pool_getter = pool_getter
self._schema = schema
self._static_schema = schema
self._schema_getter = schema_getter
@property
def _schema(self) -> str | None:
"""Resolve schema dynamically per-request when schema_getter is provided."""
if self._schema_getter:
return self._schema_getter()
return self._static_schema
async def store(
self,