Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/models/configuration.py: 62%

1from typing import Optional 

2 

3import sqlalchemy as sa 

4from sqlalchemy.ext.asyncio import AsyncSession 

5 

6from prefect.server import schemas 

7from prefect.server.database import PrefectDBInterface, db_injector, orm_models 

8 

9 

10@db_injector 

11async def write_configuration( 

12 db: PrefectDBInterface, 

13 session: AsyncSession, 

14 configuration: schemas.core.Configuration, 

15) -> orm_models.Configuration: 

16 # first see if the key already exists 

17 query = sa.select(db.Configuration).where(db.Configuration.key == configuration.key) 

18 result = await session.execute(query) # type: ignore 

19 existing_configuration = result.scalar() 

20 # if it exists, update its value 

21 if existing_configuration: 

22 existing_configuration.value = configuration.value 

23 # else create a new ORM object 

24 else: 

25 existing_configuration = db.Configuration( 

26 key=configuration.key, value=configuration.value 

27 ) 

28 session.add(existing_configuration) 

29 await session.flush() 

30 

31 # clear the cache for this key after writing a value 

32 db.queries.clear_configuration_value_cache_for_key(key=configuration.key) 

33 

34 return existing_configuration 

35 

36 

37@db_injector 

38async def delete_configuration( 

39 db: PrefectDBInterface, 

40 session: AsyncSession, 

41 key: str, 

42) -> bool: 

43 query = sa.delete(db.Configuration).where(db.Configuration.key == key) 

44 result = await session.execute(query) 

45 db.queries.clear_configuration_value_cache_for_key(key=key) 

46 return result.rowcount > 0 

47 

48 

49@db_injector 

50async def read_configuration( 

51 db: PrefectDBInterface, 

52 session: AsyncSession, 

53 key: str, 

54) -> Optional[schemas.core.Configuration]: 

55 value = await db.queries.read_configuration_value(session=session, key=key) 

56 return ( 

57 schemas.core.Configuration(key=key, value=value) if value is not None else None 

58 )