Fix CREATE TYPE race that can crash the worker on first boot

CHANGELOG.md

@@ -37,6 +37,7 @@ For more information about each release including git tags and artifacts, see [R
 
 ### Fixed
 
+- Serialize Postgres schema creation behind an advisory lock so the worker and UI API no longer race on `CREATE TYPE` for the `job_status` enum on a fresh volume, which could crash the worker on first boot ([#436](https://github.com/roostorg/osprey/pull/436) by [@vedarolap](https://github.com/vedarolap), closes [#432](https://github.com/roostorg/osprey/issues/432))
 - Escape literal braces when parsing f-strings in the engine ([#347](https://github.com/roostorg/osprey/pull/347) by [@haileyok](https://github.com/haileyok))
 - Tolerate malformed URI escapes in `EntityWithPopover` UI component ([#377](https://github.com/roostorg/osprey/pull/377) by [@julietshen](https://github.com/julietshen))
 - Add retention limits to Kafka topics to prevent unbounded disk growth ([#249](https://github.com/roostorg/osprey/pull/249) by [@VINODvoid](https://github.com/VINODvoid))

osprey_worker/src/osprey/worker/lib/storage/postgres.py

@@ -19,6 +19,10 @@
 metadata = MetaData()
 Model = declarative_base(name='Model', metadata=metadata)
 
+# Arbitrary key for the Postgres advisory lock guarding schema creation (see init_from_config).
+# Any process taking this lock is guaranteed to be alone while it runs metadata.create_all().
+_SCHEMA_CREATE_LOCK_KEY = 87271
+
 if TYPE_CHECKING:
     SessionMaker = sessionmaker[Session]  # type: ignore[type-var]
 else:
@@ -36,6 +40,24 @@ def _get_or_init_session(database: str) -> SessionMaker:
     return sessions[database]
 
 
+def create_schema(engine: sqlalchemy.engine.Engine) -> None:
+    """Create all tables/types defined in `metadata` against `engine`.
+
+    Multiple processes (e.g. the worker and the UI API) can call this at the same time against
+    a fresh database. SQLAlchemy's enum creation is check-then-create rather than atomic, so on
+    a fresh volume both processes can see a type as missing and both issue CREATE TYPE, and the
+    loser crashes with a UniqueViolation. Take a Postgres advisory lock first so only one process
+    creates the schema at a time; by the time any other process acquires the lock, create_all's
+    own existence checks make it a no-op.
+    """
+    with engine.connect() as connection:
+        connection.execute(sqlalchemy.text('SELECT pg_advisory_lock(:key)'), {'key': _SCHEMA_CREATE_LOCK_KEY})
+        try:
+            metadata.create_all(engine)
+        finally:
+            connection.execute(sqlalchemy.text('SELECT pg_advisory_unlock(:key)'), {'key': _SCHEMA_CREATE_LOCK_KEY})
+
+
 def init_from_config(database: str) -> None:
     def _init(config: Config) -> None:
         if not config['POSTGRES_HOSTS'].get(database):
@@ -59,8 +81,7 @@ def _init(config: Config) -> None:
             temporary_ability_token,
         )
 
-        # Create all tables defined in the metadata
-        metadata.create_all(new_engine)
+        create_schema(new_engine)
 
     CONFIG.instance().register_configuration_callback(_init)
 

osprey_worker/src/osprey/worker/lib/storage/tests/test_postgres.py

@@ -0,0 +1,57 @@
+import threading
+
+import sqlalchemy
+from osprey.worker.lib.singletons import CONFIG
+from osprey.worker.lib.storage import postgres
+from psycopg2.errors import DuplicateDatabase, InvalidCatalogName
+from sqlalchemy.exc import ProgrammingError
+from sqlalchemy_utils import create_database, drop_database
+
+
+def test_create_schema_survives_concurrent_callers_on_a_fresh_database():
+    """Regression test for issue #432: on a fresh database, two processes (e.g. the worker and
+    the UI API) both calling create_schema() at startup used to be able to race on `CREATE TYPE`
+    for the job_status enum, since SQLAlchemy's enum creation is check-then-create rather than
+    atomic. Both would see the type as missing, both would issue CREATE TYPE, and the loser would
+    crash with a UniqueViolation on `pg_type_typname_nsp_index`.
+
+    Simulate two racing processes with two independent engines hitting a brand new database at
+    the same time, and assert neither raises.
+    """
+    base_url = CONFIG.instance()['POSTGRES_HOSTS']['osprey_db']
+    fresh_url = base_url.rsplit('/', 1)[0] + '/osprey_test_concurrent_schema_create'
+
+    try:
+        drop_database(fresh_url)
+    except ProgrammingError as e:
+        if not isinstance(e.orig, InvalidCatalogName):
+            raise
+    try:
+        create_database(fresh_url)
+    except ProgrammingError as e:
+        if not isinstance(e.orig, DuplicateDatabase):
+            raise
+
+    try:
+        errors: list[Exception] = []
+        barrier = threading.Barrier(2)
+
+        def _create_schema() -> None:
+            engine = sqlalchemy.create_engine(fresh_url)
+            try:
+                barrier.wait(timeout=5)
+                postgres.create_schema(engine)
+            except Exception as e:
+                errors.append(e)
+            finally:
+                engine.dispose()
+
+        threads = [threading.Thread(target=_create_schema) for _ in range(2)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join(timeout=10)
+
+        assert not errors, errors
+    finally:
+        drop_database(fresh_url)