Coverage for instanovo/scripts/set_gcp_credentials.py: 86%

1# /// script 

2# requires-python = ">=3.10" 

3# dependencies = [ 

4# "python-dotenv", 

5# ] 

6# /// 

7"""Set Google Storage credentials.""" 

8 

9from __future__ import annotations 

10 

11import base64 

12import json 

13import os 

14from pathlib import Path 

15 

16from dotenv import load_dotenv 

17 

18load_dotenv() 

19 

20 

21def to_base64(json_path: Path | str) -> str: 

22 """Convert a JSON file with GS credentials to a base64 encoded string.""" 

23 with open(json_path) as f: 

24 credentials = f.read() 

25 return base64.b64encode(credentials.encode("ascii")).decode() 

26 

27 

28def set_credentials() -> None: 

29 """Set the GS credentials. 

30 

31 - To access GCP buckets, the credentials are stored in a json file. 

32 - For runs on AIchor we only have access to the encoded string which should be decoded and saved 

33 

34 Raises: 

35 OSError: if 'GOOGLE_APPLICATION_CREDENTIALS' is not set OR 

36 if 'GOOGLE_APPLICATION_CREDENTIALS' file does not exist and 

37 'GS_CREDENTIALS_ENCODED' is not set 

38 

39 """ 

40 try: 

41 gcp_crendentials_path = Path(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) 

42 except KeyError: 

43 msg = ( 

44 "To use GCP buckets you should set 'GOOGLE_APPLICATION_CREDENTIALS' env variable. " 

45 "It corresponds to the path to the json file with the credentials." 

46 ) 

47 raise OSError(msg) from None 

48 

49 if gcp_crendentials_path.exists(): 

50 return 

51 

52 try: 

53 gcp_credentials_encoded = os.environ["GS_CREDENTIALS_ENCODED"] 

54 except KeyError: 

55 msg = ( 

56 "If the json file 'GOOGLE_APPLICATION_CREDENTIALS' does not exist, you must set 'GS_CREDENTIALS_ENCODED' as the base64 encoded json file." 

57 ) 

58 raise OSError(msg) from None 

59 

60 credentials = json.loads(base64.b64decode(gcp_credentials_encoded).decode()) 

61 with open(gcp_crendentials_path, "w") as f: 

62 json.dump(credentials, f) 

63 print(f"Created {gcp_crendentials_path}") # noqa: T201 

64 

65 

66if __name__ == "__main__": 

67 # print(to_base64('ext-dtu-denovo-sequencing-gcp-6cfd4324b948.json')) 

68 set_credentials()