Coverage for api/endpoints/orders/archive.py: 60%

1from typing import Annotated 

2 

3import aiofiles.os as aios 

4from fastapi import APIRouter, Depends 

5from handlers.authorization.check_role import has_role 

6from handlers.authorization.company_employee_access import ( 

7 flexible_company_assertion_query, 

8) 

9from handlers.authorization.confidential import hide_model_fields_from_carrier 

10from handlers.grabbers.orders import archive_data_grabber 

11from handlers.orders.archive import ( 

12 export_to_xml, 

13 get_archive_orders, 

14 on_get_archive_stats, 

15) 

16from keycloak import idp 

17from sotrans_models.models._mongo import PydanticObjectIdPath 

18from sotrans_models.models.orders.order import OrderDBModel 

19from sotrans_models.models.responses import ( 

20 ErrorRepr, 

21 GenericGetListResponse, 

22 MonthlyStatisticsResponse, 

23) 

24from sotrans_models.models.roles import SotransRole 

25from sotrans_models.models.users import SotransOIDCUserModel 

26from starlette.background import BackgroundTasks 

27from starlette.responses import FileResponse 

28from utils.access_wrapper import get_active_user 

29from utils.data_grabber import BaseGetListQueryParams, BaseGetOneQueryParams 

30from utils.excel import export_orders_to_excel 

31from utils.helper import get_org_oid 

32 

33archive_router = APIRouter(prefix="/archive", tags=["orders", "archive"]) 

34 

35 

36@archive_router.get("") 

37async def get_archive( 

38 user=Depends( 

39 get_active_user( 

40 [SotransRole.carrier_logistician, SotransRole.drivers_bot] 

41 ) 

42 ), 

43 params: BaseGetListQueryParams = Depends(), 

44) -> GenericGetListResponse[OrderDBModel]: 

45 return await get_archive_orders(user, params) 

46 

47 

48@archive_router.get("/xml") 

49async def export_archive_as_xml( 

50 user: Annotated[ 

51 SotransOIDCUserModel, 

52 Depends( 

53 idp.get_current_user( 

54 required_role_names=[SotransRole.company_logistician] 

55 ) 

56 ), 

57 ], 

58 bg_task: BackgroundTasks, 

59 params: BaseGetListQueryParams = Depends(), 

60) -> FileResponse: 

61 archive_orders = await get_archive_orders(user, params) 

62 file_path = await export_to_xml(archive_orders.items) 

63 bg_task.add_task(aios.remove, file_path) 

64 return FileResponse(file_path) 

65 

66 

67@archive_router.get("/xlsx") 

68async def export_archive_as_xlsx( 

69 user: Annotated[ 

70 SotransOIDCUserModel, 

71 Depends( 

72 idp.get_current_user( 

73 required_role_names=[SotransRole.company_logistician] 

74 ) 

75 ), 

76 ], 

77 bg_task: BackgroundTasks, 

78 params: BaseGetListQueryParams = Depends(), 

79) -> FileResponse: 

80 archive_orders = await get_archive_orders(user, params) 

81 to_cut = False if has_role(user, SotransRole.company_logistician) else True 

82 file_path = await export_orders_to_excel(archive_orders.items, to_cut) 

83 bg_task.add_task(aios.remove, file_path) 

84 return FileResponse(file_path) 

85 

86 

87@archive_router.get("/statistics") 

88async def get_archive_statistics( 

89 _: Annotated[ 

90 SotransOIDCUserModel, 

91 Depends( 

92 idp.get_current_user( 

93 required_role_names=[SotransRole.company_logistician] 

94 ) 

95 ), 

96 ], 

97) -> MonthlyStatisticsResponse: 

98 """Each number in key represents how far ago starting period was.""" 

99 return await on_get_archive_stats() 

100 

101 

102@archive_router.get("/{order_id}", responses={404: {"model": ErrorRepr}}) 

103async def get_archive_by_id( 

104 order_id: PydanticObjectIdPath, 

105 user=Depends( 

106 get_active_user( 

107 [ 

108 SotransRole.carrier_logistician, 

109 SotransRole.drivers_bot, 

110 ] 

111 ) 

112 ), 

113 params: BaseGetOneQueryParams = Depends(), 

114) -> OrderDBModel: 

115 if has_role(user, SotransRole.drivers_bot): 

116 return await archive_data_grabber.get_one(order_id, params) 

117 if has_role(user, SotransRole.company_logistician): 

118 restriction_q = await flexible_company_assertion_query(user) 

119 return await archive_data_grabber.get_one_by_id_with_pattern( 

120 order_id, params, restriction_q 

121 ) 

122 return await archive_data_grabber.get_one_by_id_with_pattern( 

123 order_id, 

124 params, 

125 {OrderDBModel.carrier.id: get_org_oid(user)}, 

126 processors=[hide_model_fields_from_carrier], 

127 )