Coverage for chalicelib/core/tenants.py: 100%

1from cachetools import TTLCache, cached 

2from decouple import config 

3 

4from chalicelib.core import license 

5from chalicelib.utils import helper 

6from chalicelib.utils import pg_client 

7 

8cache = TTLCache(maxsize=1000, ttl=config("TENANTS_CACHE_TTL_S", cast=int, default=60)) 

9 

10 

11@cached(cache) 

12def get_by_tenant_id(tenant_id): 

13 with pg_client.PostgresClient() as cur: 

14 query = cur.mogrify(f"""SELECT tenants.tenant_id, 

15 tenants.name, 

16 tenants.api_key, 

17 tenants.created_at, 

18 '{license.EDITION}' AS edition, 

19 openreplay_version() AS version_number, 

20 tenants.opt_out, 

21 scope_state 

22 FROM public.tenants 

23 LIMIT 1;""", 

24 {"tenantId": tenant_id}) 

25 cur.execute(query=query) 

26 return helper.dict_to_camel_case(cur.fetchone()) 

27 

28 

29def get_by_api_key(api_key): 

30 with pg_client.PostgresClient() as cur: 

31 query = cur.mogrify(f"""SELECT 1 AS tenant_id, 

32 tenants.name, 

33 tenants.created_at  

34 FROM public.tenants 

35 WHERE tenants.api_key = %(api_key)s 

36 LIMIT 1;""", 

37 {"api_key": api_key}) 

38 cur.execute(query=query) 

39 return helper.dict_to_camel_case(cur.fetchone()) 

40 

41 

42def generate_new_api_key(tenant_id): 

43 with pg_client.PostgresClient() as cur: 

44 query = cur.mogrify(f"""UPDATE public.tenants 

45 SET api_key=generate_api_key(20) 

46 WHERE tenant_id = %(tenant_id)s 

47 RETURNING api_key;""", 

48 {"tenant_id": tenant_id}) 

49 cur.execute(query=query) 

50 return helper.dict_to_camel_case(cur.fetchone()) 

51 

52 

53def edit_tenant(tenant_id, changes): 

54 with pg_client.PostgresClient() as cur: 

55 query = cur.mogrify(f"""UPDATE public.tenants  

56 SET {", ".join([f"{helper.key_to_snake_case(k)} = %({k})s" for k in changes.keys()])}  

57 RETURNING name, opt_out;""", 

58 {"tenant_id": tenant_id, **changes}) 

59 cur.execute(query=query) 

60 return helper.dict_to_camel_case(cur.fetchone()) 

61 

62 

63def tenants_exists_sync(use_pool=True): 

64 with pg_client.PostgresClient(use_pool=use_pool) as cur: 

65 cur.execute("SELECT EXISTS(SELECT 1 FROM public.tenants)") 

66 out = cur.fetchone()["exists"] 

67 return out 

68 

69 

70async def tenants_exists(use_pool=True): 

71 from app import app 

72 async with app.state.postgresql.connection() as cnx: 

73 async with cnx.cursor() as txn: 

74 row = await cnx.execute("SELECT EXISTS(SELECT 1 FROM public.tenants)") 

75 row = await row.fetchone() 

76 return row["exists"]