h2o-drive 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of h2o-drive might be problematic. Click here for more details.

@@ -0,0 +1,113 @@
1
+ Metadata-Version: 2.1
2
+ Name: h2o-drive
3
+ Version: 0.0.1
4
+ Summary: H2O Drive Python Client
5
+ Author: H2O.ai
6
+ Author-email: support@h2o.ai
7
+ Requires-Python: >=3.7,<4.0
8
+ Classifier: Intended Audience :: Developers
9
+ Classifier: Operating System :: OS Independent
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.7
12
+ Classifier: Programming Language :: Python :: 3.8
13
+ Classifier: Programming Language :: Python :: 3.9
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Topic :: Internet
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Requires-Dist: aioboto3 (>=9.3.1,<10.0.0)
19
+ Requires-Dist: aiobotocore (>=2.1.0,<3.0.0)
20
+ Requires-Dist: botocore (==1.23.24)
21
+ Requires-Dist: h2o-authn[discovery] (>=1.1.0,<2.0.0)
22
+ Requires-Dist: h2o-cloud-discovery (>=1.1.0,<2.0.0)
23
+ Requires-Dist: typing-extensions (>=3.7.4)
24
+ Description-Content-Type: text/markdown
25
+
26
+ # h2o-drive
27
+
28
+ H2O Drive Python Client.
29
+
30
+ ## Installation
31
+
32
+ ```sh
33
+ pip install h2o-drive
34
+ ```
35
+
36
+ ## Usage
37
+
38
+ ### Connecting to H2O Drive
39
+
40
+ ```py
41
+ import h2o_drive
42
+
43
+ drive = await h2o_drive.Drive()
44
+ ```
45
+
46
+ When used within the [H2O AI Cloud](https://h2o.ai/platform/ai-cloud/) environment or locally with the [H2O AI Cloud CLI](https://docs.h2o.ai/h2o-ai-cloud/developerguide/cli) configured, no further configuration is needed.
47
+
48
+ `h2o_drive.Drive()` can be configured via its optional parameters:
49
+
50
+ - `token`: Token or token provider to use. If not specified, it's determined automatically.
51
+ - `endpoint_url`: URL of the Drive endpoint. If not specified, it's determined automatically.
52
+ - `sts_endpoint_url`: URL of the STS endpoint. If not specified, it's determined automatically.
53
+ - `environment`: URL of the H2O.ai cloud environment. Used only when determining a missing token or URL.
54
+ - `config_path`: Path to the H2O.ai CLI config. Used only when determining a missing token or URL.
55
+ - `session_ttl_seconds`: The duration, in seconds, of the role session.
56
+ - `read_timeout_seconds`: The time in seconds until a timeout exception is thrown when reading from Drive.
57
+
58
+ ### Object Operations
59
+
60
+ To perform object operations, first navigate to the user's personal drive bucket:
61
+
62
+ ```py
63
+ my_bucket = drive.my_bucket()
64
+ ```
65
+
66
+ Continue operating at the root of the user's bucket (generally not needed) or navigate into a particular space (a.k.a. a directory path) with one of the following. The `home` space is a good place to start.
67
+
68
+ ```py
69
+ my_home_space = my_bucket.home()
70
+ # or
71
+ my_other_space = my_bucket.workspace("my_other_space")
72
+ ```
73
+
74
+ Once in a space within a user's bucket (or the root of the bucket), the following operations are possible:
75
+
76
+ - `upload_file` uploads a local file to the user's Drive. It takes two arguments:
77
+ - `file_name`: The file to upload.
78
+ - `object_name`: The name to give to the uploaded file once it becomes an object in our Drive bucket.
79
+
80
+ - `list_objects` lists objects in the user's drive. It takes one optional argument:
81
+ - `prefix`: When set, only objects whose names start with the specified prefix are returned.
82
+
83
+ - `download_file` downloads an object as a local file. It takes two arguments:
84
+ - `object_name`: The name of the object (including prefix) to download.
85
+ - `file_name`: Path to where the object should be saved as a local file.
86
+
87
+ - `delete_object` Deletes an object from the user's drive. It takes a single argument:
88
+ - `object_name`: The name of the object (including prefix) to delete.
89
+
90
+ - `generate_presigned_url` generates a presigned URL through which an object in the drive can be accessed. It takes two arguments:
91
+ - `object_name`: The name of the object (including prefix) to generate a presigned URL for.
92
+ - `ttl_seconds`: (Optional) How long, in seconds, the URL should be good for.
93
+
94
+ ## Examples
95
+
96
+ ### Example: Connect to the `home` space of a user's drive bucket, write a local file, and then list the contents
97
+
98
+ ```py
99
+ import h2o_drive
100
+
101
+ drive = await h2o_drive.Drive()
102
+
103
+ my_home_space = drive.my_bucket().home()
104
+
105
+ with open("test-file.txt", "w") as f:
106
+ f.write("Hello, world!")
107
+
108
+ await my_home_space.upload_file("test-file.txt", "uploaded-object-name.txt")
109
+
110
+ objects = await my_home_space.list_objects()
111
+ print(objects)
112
+ ```
113
+
@@ -0,0 +1,87 @@
1
+ # h2o-drive
2
+
3
+ H2O Drive Python Client.
4
+
5
+ ## Installation
6
+
7
+ ```sh
8
+ pip install h2o-drive
9
+ ```
10
+
11
+ ## Usage
12
+
13
+ ### Connecting to H2O Drive
14
+
15
+ ```py
16
+ import h2o_drive
17
+
18
+ drive = await h2o_drive.Drive()
19
+ ```
20
+
21
+ When used within the [H2O AI Cloud](https://h2o.ai/platform/ai-cloud/) environment or locally with the [H2O AI Cloud CLI](https://docs.h2o.ai/h2o-ai-cloud/developerguide/cli) configured, no further configuration is needed.
22
+
23
+ `h2o_drive.Drive()` can be configured via its optional parameters:
24
+
25
+ - `token`: Token or token provider to use. If not specified, it's determined automatically.
26
+ - `endpoint_url`: URL of the Drive endpoint. If not specified, it's determined automatically.
27
+ - `sts_endpoint_url`: URL of the STS endpoint. If not specified, it's determined automatically.
28
+ - `environment`: URL of the H2O.ai cloud environment. Used only when determining a missing token or URL.
29
+ - `config_path`: Path to the H2O.ai CLI config. Used only when determining a missing token or URL.
30
+ - `session_ttl_seconds`: The duration, in seconds, of the role session.
31
+ - `read_timeout_seconds`: The time in seconds until a timeout exception is thrown when reading from Drive.
32
+
33
+ ### Object Operations
34
+
35
+ To perform object operations, first navigate to the user's personal drive bucket:
36
+
37
+ ```py
38
+ my_bucket = drive.my_bucket()
39
+ ```
40
+
41
+ Continue operating at the root of the user's bucket (generally not needed) or navigate into a particular space (a.k.a. a directory path) with one of the following. The `home` space is a good place to start.
42
+
43
+ ```py
44
+ my_home_space = my_bucket.home()
45
+ # or
46
+ my_other_space = my_bucket.workspace("my_other_space")
47
+ ```
48
+
49
+ Once in a space within a user's bucket (or the root of the bucket), the following operations are possible:
50
+
51
+ - `upload_file` uploads a local file to the user's Drive. It takes two arguments:
52
+ - `file_name`: The file to upload.
53
+ - `object_name`: The name to give to the uploaded file once it becomes an object in our Drive bucket.
54
+
55
+ - `list_objects` lists objects in the user's drive. It takes one optional argument:
56
+ - `prefix`: When set, only objects whose names start with the specified prefix are returned.
57
+
58
+ - `download_file` downloads an object as a local file. It takes two arguments:
59
+ - `object_name`: The name of the object (including prefix) to download.
60
+ - `file_name`: Path to where the object should be saved as a local file.
61
+
62
+ - `delete_object` Deletes an object from the user's drive. It takes a single argument:
63
+ - `object_name`: The name of the object (including prefix) to delete.
64
+
65
+ - `generate_presigned_url` generates a presigned URL through which an object in the drive can be accessed. It takes two arguments:
66
+ - `object_name`: The name of the object (including prefix) to generate a presigned URL for.
67
+ - `ttl_seconds`: (Optional) How long, in seconds, the URL should be good for.
68
+
69
+ ## Examples
70
+
71
+ ### Example: Connect to the `home` space of a user's drive bucket, write a local file, and then list the contents
72
+
73
+ ```py
74
+ import h2o_drive
75
+
76
+ drive = await h2o_drive.Drive()
77
+
78
+ my_home_space = drive.my_bucket().home()
79
+
80
+ with open("test-file.txt", "w") as f:
81
+ f.write("Hello, world!")
82
+
83
+ await my_home_space.upload_file("test-file.txt", "uploaded-object-name.txt")
84
+
85
+ objects = await my_home_space.list_objects()
86
+ print(objects)
87
+ ```
@@ -0,0 +1,245 @@
1
+ import os
2
+ from typing import Optional
3
+ from typing import Union
4
+
5
+ import h2o_authn.discovery
6
+ import h2o_discovery
7
+ from h2o_drive import auth
8
+ from h2o_drive import core
9
+
10
+ CLOUD_DISCOVERY_SERVICE_NAME = "drive"
11
+
12
+ Workspace = core.Workspace
13
+ MY_BUCKET = core.MY_BUCKET
14
+
15
+
16
+ async def MyHome(
17
+ token: Union[auth.TokenProvider, str],
18
+ *,
19
+ endpoint_url: Optional[str] = None,
20
+ sts_endpoint_url: Optional[str] = None,
21
+ environment: Optional[str] = None,
22
+ config_path: Optional[str] = None,
23
+ session_ttl_seconds: int = 3600,
24
+ read_timeout_seconds: int = 60,
25
+ ) -> core.PrefixClient:
26
+ """Returns new client connected to the user's personal bucket and the HOME
27
+ workspace.
28
+
29
+ When the personal bucket does not exist, it's created.
30
+
31
+ When not specified, certain parameters are determined automatically by
32
+ either environment variables or via the H2O Cloud Discovery Client. While
33
+ the Cloud Discovery Client may apply multiple strategies, the following
34
+ rules generally apply:
35
+ - environment, if specified, will be used instead of searching for that
36
+ value in config_path.
37
+ - The H2O Cloud Discovery Client will fall back to common H2O environment
38
+ variables and configuration paths, but will only supersede values
39
+ specified here if a more specific environment variable,
40
+ H2O_CLOUD_DISCOVERY, is found. This value is often set in H2O cloud
41
+ applications.
42
+
43
+ Args:
44
+ token: Token or token provider to use. If not specified, it's determined
45
+ automatically.
46
+ endpoint_url: URL of the Drive endpoint. If not specified, it's
47
+ determined automatically.
48
+ sts_endpoint_url: URL of the STS endpoint. If not specified, it's
49
+ determined automatically.
50
+ environment: URL of the H2O.ai cloud environment. Used only when
51
+ determining a missing token or URL.
52
+ config_path: Path to the H2O.ai CLI config. Used only when determining a
53
+ missing token or URL.
54
+ session_ttl_seconds: The duration, in seconds, of the role session.
55
+ read_timeout_seconds: The time in seconds until a timeout exception is
56
+ thrown when reading from Drive.
57
+
58
+ Returns:
59
+ A Drive client connected to the user's personal bucket and the HOME
60
+ workspace.
61
+ """
62
+
63
+ my_bucket = await MyBucket(
64
+ token=token,
65
+ endpoint_url=endpoint_url,
66
+ sts_endpoint_url=sts_endpoint_url,
67
+ environment=environment,
68
+ config_path=config_path,
69
+ session_ttl_seconds=session_ttl_seconds,
70
+ read_timeout_seconds=read_timeout_seconds,
71
+ )
72
+ return my_bucket.home()
73
+
74
+
75
+ async def MyBucket(
76
+ token: Union[auth.TokenProvider, str],
77
+ *,
78
+ endpoint_url: Optional[str] = None,
79
+ sts_endpoint_url: Optional[str] = None,
80
+ environment: Optional[str] = None,
81
+ config_path: Optional[str] = None,
82
+ session_ttl_seconds: int = 3600,
83
+ read_timeout_seconds: int = 60,
84
+ ) -> core.Bucket:
85
+ """Returns new client connected to the user's personal bucket.
86
+
87
+ When the personal bucket does not exist, it's created.
88
+
89
+ When not specified, certain parameters are determined automatically by
90
+ either environment variables or via the H2O Cloud Discovery Client. While
91
+ the Cloud Discovery Client may apply multiple strategies, the following
92
+ rules generally apply:
93
+ - environment, if specified, will be used instead of searching for that
94
+ value in config_path.
95
+ - The H2O Cloud Discovery Client will fall back to common H2O environment
96
+ variables and configuration paths, but will only supersede values
97
+ specified here if a more specific environment variable,
98
+ H2O_CLOUD_DISCOVERY, is found. This value is often set in H2O cloud
99
+ applications.
100
+
101
+ Args:
102
+ token: Token or token provider to use. If not specified, it's determined
103
+ automatically.
104
+ endpoint_url: URL of the Drive endpoint. If not specified, it's
105
+ determined automatically.
106
+ sts_endpoint_url: URL of the STS endpoint. If not specified, it's
107
+ determined automatically.
108
+ environment: URL of the H2O.ai cloud environment. Used only when
109
+ determining a missing token or URL.
110
+ config_path: Path to the H2O.ai CLI config. Used only when determining a
111
+ missing token or URL.
112
+ session_ttl_seconds: The duration, in seconds, of the role session.
113
+ read_timeout_seconds: The time in seconds until a timeout exception is
114
+ thrown when reading from Drive.
115
+
116
+ Returns:
117
+ A Drive client connected to the user's personal bucket.
118
+ """
119
+ client = await Drive(
120
+ token=token,
121
+ endpoint_url=endpoint_url,
122
+ sts_endpoint_url=sts_endpoint_url,
123
+ environment=environment,
124
+ config_path=config_path,
125
+ session_ttl_seconds=session_ttl_seconds,
126
+ read_timeout_seconds=read_timeout_seconds,
127
+ )
128
+ bucket = client.my_bucket()
129
+ await bucket.ensure_created()
130
+ return bucket
131
+
132
+
133
+ async def Drive(
134
+ token: Union[auth.TokenProvider, str, None] = None,
135
+ *,
136
+ endpoint_url: Optional[str] = None,
137
+ sts_endpoint_url: Optional[str] = None,
138
+ environment: Optional[str] = None,
139
+ config_path: Optional[str] = None,
140
+ session_ttl_seconds: int = 3600,
141
+ read_timeout_seconds: int = 60,
142
+ ) -> core.Client:
143
+ """Returns a new instance of the Drive client.
144
+
145
+ When not specified, certain parameters are determined automatically by
146
+ either environment variables or via the H2O Cloud Discovery Client. While
147
+ the Cloud Discovery Client may apply multiple strategies, the following
148
+ rules generally apply:
149
+ - environment, if specified, will be used instead of searching for that
150
+ value in config_path.
151
+ - The H2O Cloud Discovery Client will fall back to common H2O environment
152
+ variables and configuration paths, but will only supersede values
153
+ specified here if a more specific environment variable,
154
+ H2O_CLOUD_DISCOVERY, is found. This value is often set in H2O cloud
155
+ applications.
156
+
157
+ Args:
158
+ token: Token or token provider to use. If not specified, it's determined
159
+ automatically.
160
+ endpoint_url: URL of the Drive endpoint. If not specified, it's
161
+ determined automatically.
162
+ sts_endpoint_url: URL of the STS endpoint. If not specified, it's
163
+ determined automatically.
164
+ environment: URL of the H2O.ai cloud environment. Used only when
165
+ determining a missing token or URL.
166
+ config_path: Path to the H2O.ai CLI config. Used only when determining a
167
+ missing token or URL.
168
+ session_ttl_seconds: The duration, in seconds, of the role session.
169
+ read_timeout_seconds: The time in seconds until a timeout exception is
170
+ thrown when reading from Drive.
171
+
172
+ Returns:
173
+ A Drive client.
174
+ """
175
+ if token is None or endpoint_url is None:
176
+ discovery = h2o_discovery.discover(
177
+ environment=environment, config_path=config_path
178
+ )
179
+
180
+ if token is None:
181
+ token_provider = discover_token_provider(discovery)
182
+ elif isinstance(token, auth.TokenProvider):
183
+ token_provider = token
184
+ else:
185
+ token_provider = auth.StaticTokenProvider(token)
186
+
187
+ if endpoint_url is None:
188
+ endpoint_url = os.getenv("H2O_DRIVE_ENDPOINT")
189
+ if endpoint_url is None:
190
+ service = discovery.services.get(CLOUD_DISCOVERY_SERVICE_NAME)
191
+ if service is not None:
192
+ endpoint_url = service.uri
193
+ if endpoint_url is None:
194
+ raise LookupError(
195
+ "Failed to automatically determine the endpoint URL. "
196
+ "Drive entry missing from Cloud Discovery service."
197
+ )
198
+ endpoint_url = service.uri
199
+
200
+ if sts_endpoint_url is None:
201
+ sts_endpoint_url = os.getenv("H2O_DRIVE_STS_ENDPOINT")
202
+ if sts_endpoint_url is None:
203
+ sts_endpoint_url = f'{endpoint_url.rstrip("/")}/sts/'
204
+
205
+ return await connect(
206
+ token_provider=token_provider,
207
+ sts_endpoint_url=sts_endpoint_url,
208
+ endpoint_url=endpoint_url,
209
+ session_ttl_seconds=session_ttl_seconds,
210
+ read_timeout_seconds=read_timeout_seconds,
211
+ )
212
+
213
+
214
+ async def connect(
215
+ token_provider: auth.TokenProvider,
216
+ *,
217
+ endpoint_url: str,
218
+ sts_endpoint_url: str,
219
+ session_ttl_seconds: int = 3600,
220
+ read_timeout_seconds: int = 60,
221
+ ) -> core.Client:
222
+ """Returns a new instance of the Client with session created against the
223
+ token returned by the token provider.
224
+ """
225
+
226
+ session_provider = auth.RefreshingWebIdentitySessionProvider(
227
+ token_provider=token_provider,
228
+ sts_endpoint_url=sts_endpoint_url,
229
+ session_ttl_seconds=session_ttl_seconds,
230
+ read_timeout_seconds=read_timeout_seconds,
231
+ )
232
+ return core.Client(session=session_provider, endpoint_url=endpoint_url)
233
+
234
+
235
+ def discover_token_provider(discovery: h2o_discovery.Discovery) -> auth.TokenProvider:
236
+ """Returns a new instance of an async TokenProvider configured from the
237
+ given Discovery object.
238
+ """
239
+ try:
240
+ return h2o_authn.discovery.create_async(discovery)
241
+ except KeyError:
242
+ raise LookupError(
243
+ "Failed to automatically determine the token provider from "
244
+ "information available in the environment."
245
+ )
@@ -0,0 +1,123 @@
1
+ import datetime
2
+ from typing import Optional
3
+
4
+ import aioboto3
5
+ import aiobotocore.config
6
+ import aiobotocore.session
7
+ import typing_extensions
8
+ from typing_extensions import Protocol
9
+
10
+
11
+ DEFAULT_SESSION_EXPIRATION_THRESHOLD = datetime.timedelta(minutes=3)
12
+
13
+
14
+ @typing_extensions.runtime_checkable
15
+ class TokenProvider(Protocol):
16
+ """Defines the interface of the token provider.
17
+
18
+ Token provider is used to obtain access token that is exchanged for the session
19
+ token.
20
+ """
21
+
22
+ async def __call__(self) -> str:
23
+ """Returns fresh access token for the access to H2O Drive server."""
24
+
25
+
26
+ class StaticTokenProvider:
27
+ """Implements TokenProvider interface and always returns static token passed."""
28
+
29
+ def __init__(self, token: str) -> None:
30
+ self._token = token
31
+
32
+ async def __call__(self) -> str:
33
+ return self._token
34
+
35
+
36
+ @typing_extensions.runtime_checkable
37
+ class SessionProvider(Protocol):
38
+ """Defines the interface of the session provider.
39
+
40
+ Session provider is used by the core components to get access to the H2O Drive
41
+ server.
42
+ """
43
+
44
+ async def __call__(self) -> aioboto3.Session:
45
+ pass
46
+
47
+
48
+ class StaticSessionProvider:
49
+ """Implements SessionProvider interface and always returns the session passed
50
+ to the constructor.
51
+ """
52
+
53
+ def __init__(self, session: aioboto3.Session) -> None:
54
+ self._session = session
55
+
56
+ async def __call__(self) -> aioboto3.Session:
57
+ return self._session
58
+
59
+
60
+ class RefreshingWebIdentitySessionProvider:
61
+ """Uses token provider passed to the constructor to obtain fresh session credentials
62
+ when needed.
63
+ """
64
+
65
+ def __init__(
66
+ self,
67
+ token_provider: TokenProvider,
68
+ sts_endpoint_url: str,
69
+ *,
70
+ session_ttl_seconds: int = 3600,
71
+ expiration_threshold: datetime.timedelta = DEFAULT_SESSION_EXPIRATION_THRESHOLD,
72
+ read_timeout_seconds: int = 60,
73
+ ) -> None:
74
+ self._token_provider = token_provider
75
+ self._sts_endpoint_url = sts_endpoint_url
76
+ self._expiration_threshold = expiration_threshold
77
+ self._session_ttl_seconds = session_ttl_seconds
78
+ self._read_timeout_seconds = read_timeout_seconds
79
+
80
+ self._session: Optional[aioboto3.Session] = None
81
+ self._credentials_expiration: Optional[datetime.datetime] = None
82
+
83
+ async def __call__(self) -> aioboto3.Session:
84
+ if self.refresh_required():
85
+ await self._do_refresh()
86
+ return self._session
87
+
88
+ def refresh_required(self) -> bool:
89
+ """Returns True when there's no access token set or the current one requires
90
+ refresh.
91
+ """
92
+ if self._session is None or self._credentials_expiration is None:
93
+ return True
94
+
95
+ now = datetime.datetime.now(datetime.timezone.utc)
96
+ return self._credentials_expiration <= (now + self._expiration_threshold)
97
+
98
+ async def _do_refresh(self) -> aioboto3.Session:
99
+ credentials = await self._get_credentials()
100
+ self._credentials_expiration = credentials.get("Expiration")
101
+ self._session = aioboto3.Session(
102
+ aws_access_key_id=credentials.get("AccessKeyId"),
103
+ aws_secret_access_key=credentials.get("SecretAccessKey"),
104
+ aws_session_token=credentials.get("SessionToken"),
105
+ )
106
+
107
+ async def _get_credentials(self):
108
+ boto_session = aiobotocore.session.get_session()
109
+ async with boto_session.create_client(
110
+ "sts",
111
+ endpoint_url=self._sts_endpoint_url,
112
+ config=aiobotocore.config.AioConfig(
113
+ read_timeout=self._read_timeout_seconds
114
+ ),
115
+ ) as sts:
116
+ assume_role = await sts.assume_role_with_web_identity(
117
+ # RoleArn is an arbitrary string from 20 to 2048 characters long.
118
+ RoleArn="h2odriveh2odriveh2odrive",
119
+ RoleSessionName="h2odrive",
120
+ WebIdentityToken=await self._token_provider(),
121
+ DurationSeconds=self._session_ttl_seconds,
122
+ )
123
+ return assume_role.get("Credentials", dict())
@@ -0,0 +1,315 @@
1
+ import contextlib
2
+ import datetime
3
+ import enum
4
+ import os
5
+ import pathlib
6
+ from typing import List
7
+ from typing import NamedTuple
8
+ from typing import Optional
9
+
10
+ import boto3.s3.transfer
11
+ import botocore.exceptions
12
+
13
+ from h2o_drive import auth
14
+
15
+ MY_BUCKET = "_my"
16
+ _HOME_PREFIX = "home/"
17
+ _DEFAULT_PRESIGNED_URL_TTL_SECONDS = 24 * 60 * 60 # 1 day.
18
+
19
+
20
+ class Workspace(enum.Enum):
21
+ """Defines sopported predefined workspaces."""
22
+
23
+ HOME = enum.auto()
24
+ APP = enum.auto()
25
+ APP_VERSION = enum.auto()
26
+ APP_INSTANCE = enum.auto()
27
+
28
+
29
+ class PresignedURL(str):
30
+ """Defines presigned URL as it's own type so it can be extended for integration with
31
+ other service clients (e.g. H2O-3)
32
+ """
33
+
34
+
35
+ class ObjectSummary(NamedTuple):
36
+ """Represents basic information about object."""
37
+
38
+ key: str
39
+ size: int
40
+ last_modified: datetime.datetime
41
+
42
+
43
+ class PrefixClient:
44
+ """Implements standard object operations but all keys are automatically prefixed
45
+ with configured prefix.
46
+
47
+ Simulates behavior of the bucket within the bucket.
48
+ """
49
+
50
+ def __init__(self, *, bucket: "Bucket", prefix: str) -> None:
51
+ self._bucket = bucket
52
+ self._prefix = prefix
53
+ self._prefix_length = len(prefix)
54
+
55
+ async def upload_file(self, file_name: os.PathLike, object_name: str):
56
+ """Uploads local file to the H2O Drive as object."""
57
+ await self._bucket.upload_file(file_name, self._key_with_prefix(object_name))
58
+
59
+ async def download_file(self, object_name: str, file_name: os.PathLike):
60
+ """Download H2O Drive object to local file."""
61
+ await self._bucket.download_file(self._key_with_prefix(object_name), file_name)
62
+
63
+ async def list_objects(self, prefix: Optional[str] = None) -> List[ObjectSummary]:
64
+ """Returns list of available objects.
65
+
66
+ Only the objects with the prefix associated with the instance are returned and
67
+ the prefix is stripped. Objects can be additionally filtered by further prefix.
68
+ """
69
+ filter_prefix = self._prefix
70
+ if prefix is not None:
71
+ filter_prefix = self._key_with_prefix(prefix)
72
+
73
+ objects = await self._bucket.list_objects(prefix=filter_prefix)
74
+
75
+ result: List[ObjectSummary] = []
76
+ for o in objects:
77
+ obj_name = o.key
78
+ if obj_name.startswith(self._prefix):
79
+ obj_name = obj_name[self._prefix_length :]
80
+ result.append(
81
+ ObjectSummary(key=obj_name, size=o.size, last_modified=o.last_modified)
82
+ )
83
+
84
+ return result
85
+
86
+ async def delete_object(self, object_name: str):
87
+ """Deletes H2O Drive object."""
88
+ await self._bucket.delete_object(self._key_with_prefix(object_name))
89
+
90
+ async def generate_presigned_url(
91
+ self, object_name: str, *, ttl_seconds: int = _DEFAULT_PRESIGNED_URL_TTL_SECONDS
92
+ ) -> PresignedURL:
93
+ """Generates pre-signed URL for the H2O Drive object.
94
+
95
+ Generated URL might not be accessible externally.
96
+ """
97
+ return await self._bucket.generate_presigned_url(
98
+ object_name=self._key_with_prefix(object_name), ttl_seconds=ttl_seconds
99
+ )
100
+
101
+ def _key_with_prefix(self, key: str) -> str:
102
+ return str(pathlib.Path(self._prefix) / key)
103
+
104
+
105
+ class Bucket:
106
+ def __init__(
107
+ self, name: str, *, session: auth.SessionProvider, endpoint_url: str
108
+ ) -> None:
109
+ self._name = name
110
+ self._session = session
111
+ self._endpoint_url = endpoint_url
112
+
113
+ async def upload_file(self, file_name: os.PathLike, object_name: str):
114
+ """Uploads local file to the H2O Drive as object in the bucket."""
115
+ async with self._s3_client() as s3:
116
+ await s3.upload_file(
117
+ file_name,
118
+ self._name,
119
+ object_name,
120
+ Config=boto3.s3.transfer.TransferConfig(max_io_queue=5),
121
+ )
122
+
123
+ async def download_file(self, object_name: str, file_name: os.PathLike):
124
+ """Download H2O Drive object in the bucket to local file."""
125
+ async with self._s3_client() as s3:
126
+ await s3.download_file(self._name, object_name, file_name)
127
+
128
+ async def list_objects(self, prefix: Optional[str] = None) -> List[ObjectSummary]:
129
+ """Returns list of objects available in the bucket.
130
+
131
+ Files can be filtered by the prefix.
132
+ """
133
+ async with self._s3_resource() as s3:
134
+ bucket = await s3.Bucket(self._name)
135
+ if prefix is not None:
136
+ objects = bucket.objects.filter(Prefix=prefix)
137
+ else:
138
+ objects = bucket.objects.all()
139
+ result: List[ObjectSummary] = []
140
+ async for o in objects:
141
+ result.append(
142
+ ObjectSummary(
143
+ key=o.key,
144
+ size=await o.size,
145
+ last_modified=await o.last_modified,
146
+ )
147
+ )
148
+ return result
149
+
150
+ async def delete_object(self, object_name: str):
151
+ """Deletes H2O Drive object in the bucket."""
152
+ async with self._s3_resource() as s3:
153
+ bucket = await s3.Bucket(self._name)
154
+ obj = await bucket.Object(object_name)
155
+ await obj.delete()
156
+
157
+ async def generate_presigned_url(
158
+ self, object_name: str, *, ttl_seconds: int = _DEFAULT_PRESIGNED_URL_TTL_SECONDS
159
+ ) -> PresignedURL:
160
+ """Generates pre-signed URL for the H2O Drive object in the bucket.
161
+
162
+ Generated URL might not be accessible externally.
163
+ """
164
+
165
+ async with self._s3_client() as s3:
166
+ url_str = await s3.generate_presigned_url(
167
+ ClientMethod="get_object",
168
+ Params={"Bucket": self._name, "Key": object_name},
169
+ ExpiresIn=ttl_seconds,
170
+ )
171
+ return PresignedURL(url_str)
172
+
173
+ async def create(self) -> "Bucket":
174
+ """Creates the bucket."""
175
+ async with self._s3_resource() as s3:
176
+ bucket = await s3.Bucket(self._name)
177
+ await bucket.create()
178
+ return self
179
+
180
+ async def ensure_created(self) -> "Bucket":
181
+ """Creates the bucket if it does not exist."""
182
+ try:
183
+ await self.create()
184
+ except botocore.exceptions.ClientError as error:
185
+ if error.response["Error"]["Code"] != "BucketAlreadyExists":
186
+ raise
187
+
188
+ return self
189
+
190
+ def with_prefix(self, prefix: str) -> PrefixClient:
191
+ """Returns a new instance of the Drive client operating within the bucket but
192
+ with all object operations prefixed with the prefix.
193
+ """
194
+ return PrefixClient(bucket=self, prefix=prefix)
195
+
196
+ def workspace(self, workspace: Workspace) -> PrefixClient:
197
+ """Returns new instance og Drive client operating withing the bucket but
198
+ with all object operations prefixed with prefix pre-defined by workspace.
199
+ """
200
+ if workspace == Workspace.HOME:
201
+ prefix = _HOME_PREFIX
202
+ elif workspace == Workspace.APP:
203
+ prefix = _get_app_workspace_prefix()
204
+ elif workspace == Workspace.APP_VERSION:
205
+ prefix = _get_app_version_workspace_prefix()
206
+ elif workspace == Workspace.APP_INSTANCE:
207
+ prefix = _get_app_instance_workspace_prefix()
208
+ else:
209
+ raise ValueError(f"Unsupported workspace: {workspace}.")
210
+
211
+ return self.with_prefix(prefix=prefix)
212
+
213
+ def home(self) -> PrefixClient:
214
+ """Returns new instance og Drive client operating withing the bucket but
215
+ with all object operations prefixed with prefix pre-defined by HOME workspace.
216
+ """
217
+ return self.workspace(workspace=Workspace.HOME)
218
+
219
+ @contextlib.asynccontextmanager
220
+ async def _s3_client(self):
221
+ session = await self._session()
222
+ async with session.client(
223
+ "s3",
224
+ endpoint_url=self._endpoint_url,
225
+ config=botocore.client.Config(signature_version="s3v4"),
226
+ ) as s3:
227
+ yield s3
228
+
229
+ @contextlib.asynccontextmanager
230
+ async def _s3_resource(self):
231
+ session = await self._session()
232
+ async with session.resource(
233
+ "s3",
234
+ endpoint_url=self._endpoint_url,
235
+ config=botocore.client.Config(signature_version="s3v4"),
236
+ ) as s3:
237
+ yield s3
238
+
239
+
240
+ class Client:
241
+ def __init__(self, session: auth.SessionProvider, endpoint_url: str) -> None:
242
+ self._session = session
243
+ self._endpoint_url = endpoint_url
244
+
245
+ def bucket(self, name) -> Bucket:
246
+ """Returns a new instance of the Bucket client."""
247
+ return Bucket(name=name, session=self._session, endpoint_url=self._endpoint_url)
248
+
249
+ def my_bucket(self) -> Bucket:
250
+ """Returns a new instance of the Bucket client for the user's personal
251
+ bucket."""
252
+ return Bucket(
253
+ name=MY_BUCKET, session=self._session, endpoint_url=self._endpoint_url
254
+ )
255
+
256
+
257
+ def _get_app_workspace_prefix():
258
+ app = _get_app_name()
259
+ return f"{app}/workspace/"
260
+
261
+
262
+ def _get_app_version_workspace_prefix():
263
+ app_version_prefix = _get_app_version_prefix()
264
+ return f"{app_version_prefix}/workspace/"
265
+
266
+
267
+ def _get_app_version_prefix():
268
+ app, version, app_id = _get_app_name(), _get_app_version(), _get_app_id()
269
+ return f"{app}/{app_id}-{version}"
270
+
271
+
272
+ def _get_app_instance_workspace_prefix():
273
+ app_version_prefix = _get_app_version_prefix()
274
+ app_instance_id = _get_app_instance_id()
275
+ return f"{app_version_prefix}/{app_instance_id}/workspace/"
276
+
277
+
278
+ def _get_app_name() -> str:
279
+ app = os.getenv("H2O_CLOUD_APP_NAME")
280
+ if not app:
281
+ raise LookupError(
282
+ "Unable to determine app name. "
283
+ "H2O_CLOUD_APP_NAME environment variable not set."
284
+ )
285
+ return app
286
+
287
+
288
+ def _get_app_version() -> str:
289
+ version = os.getenv("H2O_CLOUD_APP_VERSION")
290
+ if not version:
291
+ raise LookupError(
292
+ "Unable to determine app version. "
293
+ "H2O_CLOUD_APP_VERSION environment variable not set."
294
+ )
295
+ return version
296
+
297
+
298
+ def _get_app_id() -> str:
299
+ app_id = os.getenv("H2O_CLOUD_APP_ID")
300
+ if not app_id:
301
+ raise LookupError(
302
+ "Unable to determine app ID. "
303
+ "H2O_CLOUD_APP_ID environment variable not set."
304
+ )
305
+ return app_id
306
+
307
+
308
+ def _get_app_instance_id() -> str:
309
+ instance_id = os.getenv("H2O_CLOUD_INSTANCE_ID")
310
+ if not instance_id:
311
+ raise LookupError(
312
+ "Unable to determine app instance ID. "
313
+ "H2O_CLOUD_INSTANCE_ID environment variable not set."
314
+ )
315
+ return instance_id
@@ -0,0 +1,50 @@
1
+ [tool.poetry]
2
+ name = "h2o-drive"
3
+ version = "0.0.1"
4
+ description = "H2O Drive Python Client"
5
+ readme = "PUBLIC_README.md"
6
+ authors = ["H2O.ai <support@h2o.ai>"]
7
+ classifiers = [
8
+ "Intended Audience :: Developers",
9
+ "Operating System :: OS Independent",
10
+ "Topic :: Internet",
11
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
12
+ ]
13
+
14
+ [tool.poetry.dependencies]
15
+ python = "^3.7"
16
+ aioboto3 = "^9.3.1"
17
+ aiobotocore = "^2.1.0"
18
+ botocore = "1.23.24"
19
+ typing-extensions = ">=3.7.4"
20
+ h2o-authn = {version="^1.1.0", extras=["discovery"]}
21
+ h2o-cloud-discovery = "^1.1.0"
22
+
23
+ [tool.poetry.group.dev.dependencies]
24
+ black = "^23.3.0"
25
+ types-aiobotocore = {extras = ["suppport"], version = "^2.6.0"}
26
+ flake8 = "^5.0.0"
27
+ flake8-bugbear = "^23.3.12"
28
+ flake8-builtins = "^2.1.0"
29
+ flake8-deprecated = "^2.0"
30
+ flake8-eradicate = "^1.2.0"
31
+ flake8-mutable = "^1.2.0"
32
+ flake8-pyproject = "^1.2.3"
33
+ mypy = "^1.4"
34
+ boto3-stubs = "^1.28.64"
35
+
36
+ [tool.flake8]
37
+ max-line-length = 88
38
+ import-order-style = "appnexus"
39
+ application-import-names = ["h2o_drive"]
40
+ ignore = ['E203'] # Incompatible with black
41
+
42
+ [tool.mypy]
43
+
44
+ [[tool.mypy.overrides]]
45
+ module = "aioboto3"
46
+ ignore_missing_imports = true
47
+
48
+ [build-system]
49
+ requires = ["poetry-core"]
50
+ build-backend = "poetry.core.masonry.api"