aibs-informatics-aws-utils 0.0.2__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.
Files changed (49) hide show
  1. aibs_informatics_aws_utils-0.0.2/LICENSE +33 -0
  2. aibs_informatics_aws_utils-0.0.2/PKG-INFO +39 -0
  3. aibs_informatics_aws_utils-0.0.2/README.md +17 -0
  4. aibs_informatics_aws_utils-0.0.2/pyproject.toml +240 -0
  5. aibs_informatics_aws_utils-0.0.2/setup.cfg +4 -0
  6. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/__init__.py +1 -0
  7. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/_version.py +1 -0
  8. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/apigateway.py +35 -0
  9. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/athena.py +84 -0
  10. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/auth.py +59 -0
  11. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/batch.py +324 -0
  12. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/constants/__init__.py +1 -0
  13. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/constants/efs.py +44 -0
  14. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/constants/lambda_.py +19 -0
  15. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/constants/s3.py +15 -0
  16. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/core.py +345 -0
  17. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/data_sync/__init__.py +10 -0
  18. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/data_sync/file_system.py +341 -0
  19. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/data_sync/operations.py +352 -0
  20. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/dynamodb/__init__.py +61 -0
  21. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/dynamodb/conditions.py +298 -0
  22. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/dynamodb/functions.py +403 -0
  23. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/dynamodb/table.py +837 -0
  24. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ec2.py +581 -0
  25. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ecr/__init__.py +39 -0
  26. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ecr/core.py +856 -0
  27. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ecr/image_replicator.py +466 -0
  28. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ecs.py +35 -0
  29. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/efs/__init__.py +27 -0
  30. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/efs/core.py +240 -0
  31. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/efs/mount_point.py +555 -0
  32. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/efs/paths.py +147 -0
  33. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/exceptions.py +48 -0
  34. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/fsx.py +262 -0
  35. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/lambda_.py +109 -0
  36. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/logs.py +26 -0
  37. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/py.typed +0 -0
  38. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/s3.py +1405 -0
  39. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/secretsmanager.py +41 -0
  40. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ses.py +175 -0
  41. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/sns.py +97 -0
  42. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/sqs.py +104 -0
  43. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/ssm.py +84 -0
  44. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils/stepfn.py +292 -0
  45. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils.egg-info/PKG-INFO +39 -0
  46. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils.egg-info/SOURCES.txt +47 -0
  47. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils.egg-info/dependency_links.txt +1 -0
  48. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils.egg-info/requires.txt +12 -0
  49. aibs_informatics_aws_utils-0.0.2/src/aibs_informatics_aws_utils.egg-info/top_level.txt +1 -0
@@ -0,0 +1,33 @@
1
+ Allen Institute Software License – This software license is the 2-clause BSD
2
+ license plus a third clause that prohibits redistribution and use for
3
+ commercial purposes without further permission.
4
+
5
+ Copyright © 2024. Allen Institute. All rights reserved.
6
+
7
+ Redistribution and use in source and binary forms, with or without
8
+ modification, are permitted provided that the following conditions are met:
9
+
10
+ 1. Redistributions of source code must retain the above copyright notice, this
11
+ list of conditions and the following disclaimer.
12
+
13
+ 2. Redistributions in binary form must reproduce the above copyright notice,
14
+ this list of conditions and the following disclaimer in the documentation
15
+ and/or other materials provided with the distribution.
16
+
17
+ 3. Redistributions and use for commercial purposes are not permitted without
18
+ the Allen Institute’s written permission. For purposes of this license,
19
+ commercial purposes are the incorporation of the Allen Institute's software
20
+ into anything for which you will charge fees or other compensation or use of
21
+ the software to perform a commercial service for a third party. Contact
22
+ terms@alleninstitute.org for commercial licensing opportunities.
23
+
24
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
25
+ ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
26
+ WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
27
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
28
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
29
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
30
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
31
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
32
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
33
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,39 @@
1
+ Metadata-Version: 2.1
2
+ Name: aibs-informatics-aws-utils
3
+ Version: 0.0.2
4
+ Summary: Library of AWS utility code for informatics projects at the Allen Institute for Brain Science
5
+ Project-URL: Documentation, https://.github.io/aibs-informatics-aws-utils/
6
+ Project-URL: Homepage, https://github.com/AllenInstitute/aibs-informatics-aws-utils/
7
+ Project-URL: Issues, https://github.com/AllenInstitute/aibs-informatics-aws-utils/issues
8
+ Project-URL: Repository, https://github.com/AllenInstitute/aibs-informatics-aws-utils/
9
+ Requires-Python: >=3.9
10
+ Description-Content-Type: text/markdown
11
+ License-File: LICENSE
12
+ Requires-Dist: boto3~=1.35
13
+ Requires-Dist: aibs-informatics-core
14
+ Provides-Extra: dev
15
+ Requires-Dist: aibs-informatics-test-resources; extra == "dev"
16
+ Requires-Dist: boto3-stubs[apigateway,athena,batch,ecr,ecs,efs,essential,fsx,logs,secretsmanager,ses,sns,ssm,stepfunctions,sts]; extra == "dev"
17
+ Requires-Dist: moto[all]~=5.0; extra == "dev"
18
+ Provides-Extra: release
19
+ Requires-Dist: build; extra == "release"
20
+ Requires-Dist: bump-my-version; extra == "release"
21
+ Requires-Dist: wheel; extra == "release"
22
+
23
+ # AIBS Informatics AWS Utils
24
+
25
+ [![Build Status](https://github.com/AllenInstitute/aibs-informatics-aws-utils/actions/workflows/build.yml/badge.svg)](https://github.com/AllenInstitute/aibs-informatics-aws-utils/actions/workflows/build.yml)
26
+
27
+ ---
28
+
29
+ ## Overview
30
+
31
+ The AIBS Informatics AWS Utils library provides a collection of utilities and tools for working with AWS services. This library includes functionalities for interacting with AWS S3, ECR, Lambda, and other AWS services, making it easier to integrate AWS capabilities into various projects at the Allen Institute for Brain Science.
32
+
33
+ ## Contributing
34
+
35
+ Any and all PRs are welcome. Please see [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
36
+
37
+ ## Licensing
38
+
39
+ This software is licensed under the Allen Institute Software License, which is the 2-clause BSD license plus a third clause that prohibits redistribution and use for commercial purposes without further permission. For more information, please visit [Allen Institute Terms of Use](https://alleninstitute.org/terms-of-use/).
@@ -0,0 +1,17 @@
1
+ # AIBS Informatics AWS Utils
2
+
3
+ [![Build Status](https://github.com/AllenInstitute/aibs-informatics-aws-utils/actions/workflows/build.yml/badge.svg)](https://github.com/AllenInstitute/aibs-informatics-aws-utils/actions/workflows/build.yml)
4
+
5
+ ---
6
+
7
+ ## Overview
8
+
9
+ The AIBS Informatics AWS Utils library provides a collection of utilities and tools for working with AWS services. This library includes functionalities for interacting with AWS S3, ECR, Lambda, and other AWS services, making it easier to integrate AWS capabilities into various projects at the Allen Institute for Brain Science.
10
+
11
+ ## Contributing
12
+
13
+ Any and all PRs are welcome. Please see [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
14
+
15
+ ## Licensing
16
+
17
+ This software is licensed under the Allen Institute Software License, which is the 2-clause BSD license plus a third clause that prohibits redistribution and use for commercial purposes without further permission. For more information, please visit [Allen Institute Terms of Use](https://alleninstitute.org/terms-of-use/).
@@ -0,0 +1,240 @@
1
+ # -----------------------------------------------------------------------------
2
+ ## Build System Configurations
3
+ # https://setuptools.pypa.io/en/latest/setuptools.html#building-and-distributing-packages-with-setuptools
4
+ # -----------------------------------------------------------------------------
5
+
6
+ [build-system]
7
+ requires = ["setuptools>=61", "wheel"]
8
+ build-backend = "setuptools.build_meta"
9
+
10
+
11
+ # -----------------------------------------------------------------------------
12
+ [project]
13
+ name = "aibs-informatics-aws-utils"
14
+ description = "Library of AWS utility code for informatics projects at the Allen Institute for Brain Science"
15
+ readme = "README.md"
16
+ dynamic = [ "version"]
17
+ requires-python = ">=3.9"
18
+
19
+ dependencies = [
20
+ "boto3~=1.35",
21
+ "aibs-informatics-core",
22
+ ]
23
+
24
+ [project.optional-dependencies]
25
+ dev = [
26
+ "aibs-informatics-test-resources",
27
+ "boto3-stubs[athena,apigateway,batch,ecr,ecs,efs,essential,fsx,logs,secretsmanager,ses,sns,ssm,sts,stepfunctions]",
28
+ "moto[all] ~= 5.0",
29
+ ]
30
+
31
+ release = [
32
+ "build",
33
+ "bump-my-version",
34
+ "wheel",
35
+ ]
36
+
37
+ [tool.setuptools.packages.find]
38
+ where = ["src"]
39
+
40
+ [tool.setuptools.package-data]
41
+ "*" = ['py.typed']
42
+
43
+ [tool.setuptools.dynamic]
44
+ version = {attr = "aibs_informatics_aws_utils._version.__version__"}
45
+
46
+ [project.urls]
47
+ Documentation = "https://.github.io/aibs-informatics-aws-utils/"
48
+ Homepage = "https://github.com/AllenInstitute/aibs-informatics-aws-utils/"
49
+ Issues = "https://github.com/AllenInstitute/aibs-informatics-aws-utils/issues"
50
+ Repository = "https://github.com/AllenInstitute/aibs-informatics-aws-utils/"
51
+
52
+
53
+ # -----------------------------------------------------------------------------
54
+ ## Pyright Configurations
55
+ # https://github.com/microsoft/pyright/blob/main/docs/configuration.md
56
+ # -----------------------------------------------------------------------------
57
+
58
+ [tool.pyright]
59
+ reportGeneralTypeIssues = false
60
+ typeCheckingMode = "basic"
61
+
62
+ # -----------------------------------------------------------------------------
63
+ ## Coverage Configurations
64
+ # https://coverage.readthedocs.io/en/7.0.4/config.html#
65
+ # -----------------------------------------------------------------------------
66
+
67
+ [tool.coverage]
68
+ # Note: we use pytest-cov to generate coverage reports
69
+ # when running tests. but these coverage configs
70
+ # are read when doing so.
71
+ [tool.coverage.run]
72
+ branch = true
73
+ command_line = "coverage -m pytest"
74
+ data_file = "build/.coverage"
75
+ source = [
76
+ "src/"
77
+ ]
78
+
79
+ [tool.coverage.report]
80
+ omit = [
81
+ "test/*"
82
+ ]
83
+ exclude_lines = [
84
+ 'pragma: no cover',
85
+ 'raise NotImplementedError',
86
+ 'if TYPE_CHECKING:',
87
+ 'if typing.TYPE_CHECKING:',
88
+ '@overload',
89
+ '@typing.overload',
90
+ '\(Protocol\):$',
91
+ 'typing.assert_never',
92
+ 'assert_never',
93
+ ]
94
+ skip_empty = true
95
+
96
+ [tool.coverage.html]
97
+ directory = "build/documentation/coverage"
98
+
99
+ [tool.coverage.xml]
100
+ output = "build/documentation/coverage.xml"
101
+
102
+ # -----------------------------------------------------------------------------
103
+ ## pytest Configurations
104
+ # https://docs.pytest.org/en/7.2.x/reference/customize.html
105
+ # -----------------------------------------------------------------------------
106
+
107
+ [tool.pytest.ini_options]
108
+ minversion = "6.0"
109
+ addopts = [
110
+ "-ra",
111
+ "--verbose",
112
+ "--ignore=build/private",
113
+ # Coverage options should be managed in the .coveragerc file.
114
+ # The below configurations simply enable coverage and reporting.
115
+ "--cov",
116
+ "--cov-report=term-missing",
117
+ "--cov-report=html",
118
+ "--cov-report=xml",
119
+ "--cov-fail-under=0",
120
+ "--color=yes",
121
+ ]
122
+ testpaths = [
123
+ "test",
124
+ ]
125
+ cache_dir = "build/.pytest_cache"
126
+
127
+ # -----------------------------------------------------------------------------
128
+ ## MyPy Configurations
129
+ # https://mypy.readthedocs.io/en/stable/config_file.html#example-pyproject-toml
130
+ # -----------------------------------------------------------------------------
131
+
132
+ [tool.mypy]
133
+ exclude = '''(?x)(
134
+ assets/
135
+ | build/
136
+ | build_tools
137
+ | setup.py$
138
+ | test/
139
+ )'''
140
+
141
+ cache_dir = "build/.mypy_cache"
142
+
143
+ # Using no incremental to avoid known issue with mypy and some packages:
144
+ # https://github.com/python/mypy/issues/9852
145
+ incremental = false
146
+
147
+ # Import Discovery
148
+ # https://mypy.readthedocs.io/en/stable/config_file.html#import-discovery
149
+ ignore_missing_imports = true
150
+ follow_imports = "silent"
151
+ no_site_packages = true
152
+
153
+ # Untyped definitions and calls
154
+ # https://mypy.readthedocs.io/en/stable/config_file.html#untyped-definitions-and-calls
155
+ # TODO: enable and fix errors
156
+ check_untyped_defs = false
157
+
158
+ # Miscellaneous strictness flags
159
+ # https://mypy.readthedocs.io/en/latest/config_file.html#miscellaneous-strictness-flags
160
+ allow_redefinition = true
161
+
162
+
163
+ # Configuring Error Messages
164
+ # https://mypy.readthedocs.io/en/stable/config_file.html#configuring-error-messages
165
+ show_error_codes = true
166
+ color_output = true
167
+ pretty = true
168
+ show_absolute_path = false
169
+
170
+ # Reporting Generation
171
+ # https://mypy.readthedocs.io/en/stable/config_file.html#report-generation
172
+ # [DEBUG] UNCOMMENT FOR DEBUG PURPOSES ONLY! Type coverage impacts test
173
+ # html_report = build/documentation/mypy/
174
+ # xml_report = build/documentation/mypy/
175
+ # cobertura_xml_report = build/documentation/mypy/
176
+
177
+ # None and Optional handling
178
+ # https://mypy.readthedocs.io/en/latest/config_file.html#none-and-optional-handling
179
+ strict_optional = false
180
+
181
+ # Miscellaneous
182
+ # https://mypy.readthedocs.io/en/latest/config_file.html#miscellaneous
183
+ # [DEBUG] If you need to better understand type error, increase verbosity
184
+ verbosity = 0
185
+
186
+ [[tool.mypy.overrides]]
187
+ module = [
188
+ "src.*",
189
+ "test.*"
190
+ ]
191
+ ignore_errors = false
192
+
193
+
194
+ # -----------------------------------------------------------------------------
195
+ ## Black Configurations
196
+ # https://black.readthedocs.io/en/stable/usage_and_configuration/the_basics.html#what-on-earth-is-a-pyproject-toml-file
197
+ # -----------------------------------------------------------------------------
198
+
199
+
200
+ [tool.black]
201
+ line-length = 99
202
+ include = '\.pyi?$'
203
+
204
+ # -----------------------------------------------------------------------------
205
+ ## isort Configurations
206
+ # https://pycqa.github.io/isort/docs/configuration/config_files.html
207
+ # -----------------------------------------------------------------------------
208
+
209
+ [tool.isort]
210
+ # required for compatibility with black:
211
+ line_length = 99
212
+ profile = "black"
213
+ src_paths = ["src", "test"]
214
+
215
+ # -----------------------------------------------------------------------------
216
+ ## bumpversion Configurations
217
+ # https://callowayproject.github.io/bump-my-version/
218
+ # -----------------------------------------------------------------------------
219
+
220
+ [tool.bumpversion]
221
+ allow_dirty = false
222
+ commit = true
223
+ commit_args = ""
224
+ current_version = "0.0.2"
225
+ ignore_missing_version = false
226
+ message = "Bump version: {current_version} → {new_version}"
227
+ parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
228
+ regex = false
229
+ replace = "{new_version}"
230
+ search = "{current_version}"
231
+ serialize = ["{major}.{minor}.{patch}"]
232
+ sign_tags = false
233
+ tag = true
234
+ tag_message = "Bump version: {current_version} → {new_version}"
235
+ tag_name = "v{new_version}"
236
+
237
+ [[tool.bumpversion.files]]
238
+ filename = "src/aibs_informatics_aws_utils/_version.py"
239
+ search = "__version__ = \"{current_version}\""
240
+ replace = "__version__ = \"{new_version}\""
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1 @@
1
+ from aibs_informatics_aws_utils.core import * # type: ignore
@@ -0,0 +1 @@
1
+ __version__ = "0.0.2"
@@ -0,0 +1,35 @@
1
+ from typing import TYPE_CHECKING, List
2
+
3
+ from aibs_informatics_aws_utils.core import AWSService, get_region
4
+ from aibs_informatics_aws_utils.exceptions import ResourceNotFoundError
5
+
6
+ if TYPE_CHECKING: # pragma: no cover
7
+ from mypy_boto3_apigateway.type_defs import RestApiTypeDef
8
+ else:
9
+ RestApiTypeDef = dict
10
+
11
+ get_apigateway_client = AWSService.API_GATEWAY.get_client
12
+
13
+
14
+ def get_rest_api(api_name: str, region: str = None) -> RestApiTypeDef:
15
+ apigw = get_apigateway_client(region=region)
16
+
17
+ paginator = apigw.get_paginator("get_rest_apis")
18
+ rest_apis: List[RestApiTypeDef] = paginator.paginate(
19
+ PaginationConfig={"MaxItems": 100}
20
+ ).build_full_result()["items"]
21
+
22
+ for rest_api in rest_apis:
23
+ # In theory, only one api should be associated with env-base
24
+ if rest_api.get("name") == api_name:
25
+ return rest_api
26
+ else:
27
+ raise ResourceNotFoundError(f"Could not resolve REST Api with {api_name}")
28
+
29
+
30
+ def get_rest_api_endpoint(
31
+ rest_api: RestApiTypeDef, stage: str = "prod", region: str = None
32
+ ) -> str:
33
+ api_id = rest_api["id"] # type: ignore # mypy_boto3 TypeDict makes optional, but actually is required
34
+ region = get_region(region)
35
+ return f"https://{api_id}.execute-api.{region}.amazonaws.com/{stage}"
@@ -0,0 +1,84 @@
1
+ import logging
2
+ import re
3
+ import time
4
+ from string import Template
5
+ from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
6
+
7
+ from aibs_informatics_core.env import EnvBase
8
+ from botocore.exceptions import ClientError
9
+
10
+ from aibs_informatics_aws_utils.core import AWSService
11
+ from aibs_informatics_aws_utils.exceptions import AWSError
12
+
13
+ if TYPE_CHECKING: # pragma: no cover
14
+ from mypy_boto3_athena.type_defs import (
15
+ GetQueryExecutionInputRequestTypeDef,
16
+ GetQueryExecutionOutputTypeDef,
17
+ QueryExecutionStatusTypeDef,
18
+ QueryExecutionTypeDef,
19
+ StartQueryExecutionInputRequestTypeDef,
20
+ StartQueryExecutionOutputTypeDef,
21
+ )
22
+ else:
23
+ GetQueryExecutionInputRequestTypeDef = dict
24
+ GetQueryExecutionOutputTypeDef = dict
25
+ QueryExecutionStatusTypeDef = dict
26
+ QueryExecutionTypeDef = dict
27
+ StartQueryExecutionInputRequestTypeDef = dict
28
+ StartQueryExecutionOutputTypeDef = dict
29
+
30
+
31
+ ATHENA_QUERY_WAITER_STATUS = Literal["SUCCEEDED", "FAILED", "CANCELLED", "TIMEOUT"]
32
+
33
+ logger = logging.getLogger(__name__)
34
+
35
+ get_athena_client = AWSService.ATHENA.get_client
36
+
37
+
38
+ def start_query_execution(
39
+ query_string: str,
40
+ work_group: Optional[str] = None,
41
+ execution_parameters: Optional[List[str]] = None,
42
+ **kwargs,
43
+ ) -> StartQueryExecutionOutputTypeDef:
44
+ athena = get_athena_client()
45
+
46
+ request = StartQueryExecutionInputRequestTypeDef(QueryString=query_string)
47
+ if work_group:
48
+ request["WorkGroup"] = work_group
49
+ if execution_parameters:
50
+ request["ExecutionParameters"] = execution_parameters
51
+ request.update(kwargs)
52
+ try:
53
+ metadata = athena.start_query_execution(**request)
54
+ return metadata
55
+ except ClientError as e:
56
+ logger.error(f"Error executing : {request} {e}", exc_info=True)
57
+ raise AWSError(f"Error starting query execution: {request} {e}") from e
58
+
59
+
60
+ def get_query_execution(query_execution_id: str) -> GetQueryExecutionOutputTypeDef:
61
+ athena = get_athena_client()
62
+ try:
63
+ return athena.get_query_execution(QueryExecutionId=query_execution_id)
64
+ except Exception as e:
65
+ logger.error(f"Error executing : {query_execution_id} {e}", exc_info=True)
66
+ raise AWSError(f"Error starting query execution: {query_execution_id} {e}") from e
67
+
68
+
69
+ def query_waiter(
70
+ query_execution_id: str, timeout: int = 60
71
+ ) -> Tuple[ATHENA_QUERY_WAITER_STATUS, QueryExecutionStatusTypeDef]:
72
+ start = time.time()
73
+ logger.info(f"Polling for status of query execution: {query_execution_id}")
74
+ while True:
75
+ stats = get_query_execution(query_execution_id=query_execution_id)
76
+ logger.info(f"Query Execution Status: {stats}")
77
+ status = stats["QueryExecution"].get("Status", {})
78
+ state = status.get("State")
79
+ if state and state in ["SUCCEEDED", "FAILED", "CANCELLED"]:
80
+ return state, status
81
+ time.sleep(0.2) # 200ms
82
+ # Exit if the time waiting exceed the timeout seconds
83
+ if time.time() > start + timeout:
84
+ return "TIMEOUT", status
@@ -0,0 +1,59 @@
1
+ from typing import Optional
2
+
3
+ import requests
4
+ from botocore.auth import SigV4Auth
5
+ from botocore.awsrequest import AWSRequest
6
+ from botocore.compat import parse_qsl, urlparse
7
+ from botocore.session import Session
8
+ from requests.auth import AuthBase
9
+
10
+ from aibs_informatics_aws_utils.core import get_session
11
+
12
+
13
+ class IamAWSRequestsAuth(AuthBase):
14
+ """
15
+ IAM authorizer.
16
+
17
+ :param boto3.Session session: Optional boto3 Session object
18
+ :param str service_name: Optional AWS service name
19
+
20
+ :Example:
21
+
22
+ >>> IAMAuth()
23
+ >>> IAMAuth(boto3.Session(), 'execute-api')
24
+ """
25
+
26
+ def __init__(self, session: Optional[Session] = None, service_name: str = "execute-api"):
27
+ self.boto3_session = get_session(session)
28
+ credentials = self.boto3_session.get_credentials()
29
+ if not credentials:
30
+ raise ValueError("No AWS credentials found")
31
+
32
+ self.sigv4 = SigV4Auth(
33
+ credentials=credentials.get_frozen_credentials(),
34
+ service_name=service_name,
35
+ region_name=self.boto3_session.region_name,
36
+ )
37
+
38
+ def __call__(self, request):
39
+ # Parse request URL
40
+ url = urlparse(request.url)
41
+
42
+ # Prepare AWS request
43
+ awsrequest = AWSRequest(
44
+ method=request.method,
45
+ url=f"{url.scheme}://{url.netloc}{url.path}",
46
+ data=request.body if hasattr(request, "body") else (request.json or request.data),
47
+ params=dict(parse_qsl(url.query)),
48
+ )
49
+
50
+ # Sign request
51
+ self.sigv4.add_auth(awsrequest)
52
+
53
+ # Re-add original headers
54
+ for key, val in request.headers.items():
55
+ if key not in awsrequest.headers:
56
+ awsrequest.headers[key] = val
57
+
58
+ # Return prepared request
59
+ return awsrequest.prepare()