Coverage for /home/ubuntu/hidebound/python/hidebound/exporters/s3_exporter.py: 100%

38 statements  

« prev     ^ index     » next       coverage.py v7.5.4, created at 2024-07-05 23:50 +0000

1from typing import Any, List # noqa F401 

2 

3from io import BytesIO 

4import json 

5 

6from botocore.exceptions import ClientError 

7from schematics.types import StringType 

8import boto3 as boto 

9 

10from hidebound.exporters.exporter_base import ExporterBase, ExporterConfigBase 

11import hidebound.core.validators as vd 

12# ------------------------------------------------------------------------------ 

13 

14 

15class S3Config(ExporterConfigBase): 

16 ''' 

17 A class for validating configurations supplied to S3Exporter. 

18 

19 Attributes: 

20 name (str): Name of exporter. Must be 's3'. 

21 access_key (str): AWS access key. 

22 secret_key (str): AWS secret key. 

23 bucket (str): AWS bucket name. 

24 region (str): AWS region name. Default: us-east-1. 

25 ''' 

26 name = StringType( 

27 required=True, validators=[lambda x: vd.is_eq(x, 's3')] 

28 ) # type: StringType 

29 access_key = StringType(required=True) # type: StringType 

30 secret_key = StringType(required=True) # type: StringType 

31 bucket = StringType( 

32 required=True, validators=[vd.is_bucket_name] 

33 ) # type: StringType 

34 region = StringType( 

35 required=True, validators=[vd.is_aws_region] 

36 ) # type: StringType 

37 

38 

39class S3Exporter(ExporterBase): 

40 @staticmethod 

41 def from_config(config): 

42 # type: (dict) -> S3Exporter 

43 ''' 

44 Construct a S3Exporter from a given config. 

45 

46 Args: 

47 config (dict): Config dictionary. 

48 

49 Raises: 

50 DataError: If config is invalid. 

51 

52 Returns: 

53 S3Exporter: S3Exporter instance. 

54 ''' 

55 return S3Exporter(**config) 

56 

57 def __init__( 

58 self, 

59 access_key, 

60 secret_key, 

61 bucket, 

62 region, 

63 metadata_types=['asset', 'file', 'asset-chunk', 'file-chunk'], 

64 **kwargs, 

65 ): 

66 # type: (str, str, str, str, List[str], Any) -> None 

67 ''' 

68 Constructs a S3Exporter instances and creates a bucket with given name 

69 if it does not exist. 

70 

71 Args: 

72 access_key (str): AWS access key. 

73 secret_key (str): AWS secret key. 

74 bucket (str): AWS bucket name. 

75 region (str): AWS region. 

76 metadata_types (list, optional): List of metadata types for export. 

77 Default: [asset, file, asset-chunk, file-chunk]. 

78 

79 Raises: 

80 DataError: If config is invalid. 

81 ''' 

82 super().__init__(metadata_types=metadata_types) 

83 

84 config = dict( 

85 name='s3', 

86 access_key=access_key, 

87 secret_key=secret_key, 

88 bucket=bucket, 

89 region=region, 

90 metadata_types=metadata_types, 

91 ) 

92 S3Config(config).validate() 

93 # ---------------------------------------------------------------------- 

94 

95 session = boto.session.Session( 

96 aws_access_key_id=access_key, 

97 aws_secret_access_key=secret_key, 

98 region_name=region, 

99 ) 

100 self._bucket = session.resource('s3').Bucket(bucket) 

101 

102 try: 

103 session.resource('s3').meta.client.head_bucket(Bucket=bucket) 

104 except ClientError: 

105 self._bucket.create( 

106 CreateBucketConfiguration={'LocationConstraint': region} 

107 ) 

108 

109 def _export_content(self, metadata): 

110 # type: (dict) -> None 

111 ''' 

112 Exports metadata from single JSON file in hidebound/metadata/file. 

113 

114 Args: 

115 metadata (dict): File metadata. 

116 ''' 

117 self._bucket.upload_file( 

118 metadata['filepath'], 

119 'hidebound/content/' + metadata['filepath_relative'], 

120 ) 

121 

122 def _export_asset(self, metadata): 

123 # type: (dict) -> None 

124 ''' 

125 Exports metadata from single JSON file in hidebound/metadata/asset. 

126 

127 Args: 

128 metadata (dict): Asset metadata. 

129 ''' 

130 self._bucket.upload_fileobj( 

131 BytesIO(json.dumps(metadata, indent=4).encode('utf-8')), 

132 'hidebound/metadata/asset/' + metadata['asset_id'] + '.json', 

133 ) 

134 

135 def _export_file(self, metadata): 

136 # type: (dict) -> None 

137 ''' 

138 Exports metadata from single JSON file in hidebound/metadata/file. 

139 

140 Args: 

141 metadata (dict): File metadata. 

142 ''' 

143 self._bucket.upload_fileobj( 

144 BytesIO(json.dumps(metadata, indent=4).encode('utf-8')), 

145 'hidebound/metadata/file/' + metadata['file_id'] + '.json', 

146 ) 

147 

148 def _export_asset_chunk(self, metadata): 

149 # type: (List[dict]) -> None 

150 ''' 

151 Exports list of asset metadata to a single file in 

152 hidebound/metadata/asset-chunk. 

153 

154 Args: 

155 metadata (list[dict]): Asset metadata. 

156 ''' 

157 self._bucket.upload_fileobj( 

158 BytesIO(json.dumps(metadata).encode('utf-8')), 

159 f'hidebound/metadata/asset-chunk/hidebound-asset-chunk_{self._time}.json', 

160 ) 

161 

162 def _export_file_chunk(self, metadata): 

163 # type: (List[dict]) -> None 

164 ''' 

165 Exports list of file metadata to a single file in 

166 hidebound/metadata/file-chunk. 

167 

168 Args: 

169 metadata (list[dict]): File metadata. 

170 ''' 

171 self._bucket.upload_fileobj( 

172 BytesIO(json.dumps(metadata).encode('utf-8')), 

173 f'hidebound/metadata/file-chunk/hidebound-file-chunk_{self._time}.json', 

174 )