Coverage for src / superset_io / utils.py: 100%

1import io 

2import logging 

3import zipfile 

4from pathlib import Path 

5 

6log = logging.getLogger("superset_io") 

7 

8 

9def get_version(): 

10 from importlib import metadata 

11 

12 try: 

13 return metadata.version("superset-io") 

14 except metadata.PackageNotFoundError: 

15 return "[not found] Use `uv sync` when developing!" 

16 

17 

18def zipfile_buffer_from_folder(folder: Path | str) -> io.BytesIO: 

19 """Create a zipfile.ZipFile object from a folder path in memory. 

20 Args: 

21 folder (Path | str): Path to the folder to be zipped. 

22 Returns: 

23 zipfile.ZipFile: In-memory zipfile object. 

24 """ 

25 

26 folder = Path(folder) 

27 if not folder.exists() or not folder.is_dir(): 

28 raise ValueError(f"Not a folder: {folder}") 

29 

30 root_prefix = folder.name 

31 

32 buffer = io.BytesIO() 

33 with zipfile.ZipFile(buffer, "w", compression=zipfile.ZIP_DEFLATED) as zf: 

34 for path in folder.rglob("*"): 

35 if path.is_file(): 

36 arcname = Path(root_prefix) / path.relative_to(folder) 

37 zf.write(path, arcname.as_posix()) 

38 

39 buffer.seek(0) 

40 return buffer 

41 

42 

43def zipfile_buffer_from_zipfile(zip_path: Path | str) -> io.BytesIO: 

44 """Create an in-memory copy of an existing .zip file. 

45 Args: 

46 zip_path (Path | str): Path to the existing .zip file. 

47 Returns: 

48 io.BytesIO: In-memory buffer containing the zip file contents. 

49 """ 

50 zip_path = Path(zip_path) 

51 buffer = io.BytesIO() 

52 

53 # Read the original zip file and write its contents into the buffer 

54 with Path(zip_path).open("rb") as f: 

55 buffer.write(f.read()) 

56 

57 buffer.seek(0) # Reset buffer pointer to the beginning for reading 

58 return buffer 

59 

60 

61def validate_assets_bundle_structure(zip_buffer: io.BytesIO | bytes | Path) -> None: 

62 """ 

63 Check that a zip files structure will be accepted by /api/v1/assets/import. 

64 

65 Superset assets bundles typically look like: 

66 assets_export_<timestamp>/metadata.yaml 

67 assets_export_<timestamp>/charts/area_54.yaml 

68 assets_export_<timestamp>/dashboards/world_banks_data_1.yaml 

69 assets_export_<timestamp>/databases/examples.yaml 

70 assets_export_<timestamp>/datasets/bart_lines_7.yaml 

71 

72 IMPORTANT: metadata.yaml is not at the zip root, but under a 

73 single top-level folder. 

74 """ 

75 

76 zip_io: io.BytesIO | str 

77 if isinstance(zip_buffer, (bytes, memoryview, bytearray)): 

78 zip_io = io.BytesIO(zip_buffer) 

79 elif isinstance(zip_buffer, Path): 

80 zip_io = str(zip_buffer) 

81 else: 

82 zip_io = zip_buffer 

83 

84 with zipfile.ZipFile(zip_io, "r") as zf: 

85 names = [n for n in zf.namelist() if not n.endswith("/")] 

86 

87 try: 

88 # Find metadata.yaml anywhere 

89 metadata_candidates = [n for n in names if n.endswith("metadata.yaml")] 

90 if not metadata_candidates: 

91 raise ValueError("Missing metadata.yaml in ZIP.") 

92 

93 # Enforce the expected structure: exactly one top-level folder, 

94 # and metadata.yaml under it e.g. "assets_export_20260212T092634/metadata.yaml" 

95 roots = {n.split("/", 1)[0] for n in names if "/" in n} 

96 if len(roots) != 1: 

97 raise ValueError( 

98 "Expected exactly one top-level folder in assets bundle, but found: " 

99 f"{sorted(roots)}." 

100 ) 

101 

102 root = next(iter(roots)) 

103 expected_path = f"{root}/metadata.yaml" 

104 if expected_path not in names: 

105 raise ValueError( 

106 f"metadata.yaml not found at expected path '{expected_path}'. " 

107 f"Candidates were: {metadata_candidates}" 

108 ) 

109 except ValueError: 

110 log.debug("Invalid ZIP contents:\n" + " \n".join(names)) 

111 raise