api-foundry-query-engine 0.8.11__tar.gz → 0.8.30__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 (32) hide show
  1. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/PKG-INFO +1 -1
  2. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/adapters/gateway_adapter.py +79 -9
  3. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_query_handler.py +29 -6
  4. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_select_query_handler.py +17 -3
  5. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/lambda_handler.py +15 -17
  6. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/operation.py +9 -3
  7. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/utils/api_model.py +107 -28
  8. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/pyproject.toml +1 -1
  9. api_foundry_query_engine-0.8.11/api_foundry_query_engine/connectors/athena_connector.py +0 -120
  10. api_foundry_query_engine-0.8.11/api_foundry_query_engine/services/security_service.py +0 -105
  11. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/.gitignore +0 -0
  12. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/LICENSE +0 -0
  13. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/.pre-commit-config.yaml +0 -0
  14. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/__init__.py +0 -0
  15. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/adapters/adapter.py +0 -0
  16. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/adapters/case_change_adapter.py +0 -0
  17. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/adapters/security_adapter.py +0 -0
  18. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/connectors/connection.py +0 -0
  19. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/connectors/connection_factory.py +0 -0
  20. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/connectors/oracle_connector.py +0 -0
  21. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/connectors/postgres_connection.py +0 -0
  22. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/dao.py +0 -0
  23. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/operation_dao.py +0 -0
  24. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_custom_query_handler.py +0 -0
  25. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_delete_query_handler.py +0 -0
  26. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_insert_query_handler.py +0 -0
  27. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_subselect_query_handler.py +0 -0
  28. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/dao/sql_update_query_handler.py +0 -0
  29. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/services/service.py +0 -0
  30. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/services/transactional_service.py +0 -0
  31. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/utils/app_exception.py +0 -0
  32. {api_foundry_query_engine-0.8.11 → api_foundry_query_engine-0.8.30}/api_foundry_query_engine/utils/logger.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: api-foundry-query-engine
3
- Version: 0.8.11
3
+ Version: 0.8.30
4
4
  Summary: The AWS lambda service handler use by the `api_foundry` project is a powerful tool designed to automate the deployment of REST APIs on AWS using Lambda services to access and interact with relational databases (RDBMS). This project leverages the OpenAPI specification to define and manage the APIs
5
5
  Project-URL: Documentation, https://github.com/DanRepik/api-foundry
6
6
  Project-URL: Source, https://github.com/DanRepik/api-foundry
@@ -1,7 +1,9 @@
1
1
  import json
2
+ from typing import Any, Dict, List, Optional, Tuple
2
3
 
3
4
  from api_foundry_query_engine.adapters.adapter import Adapter
4
5
  from api_foundry_query_engine.operation import Operation
6
+ from api_foundry_query_engine.utils.app_exception import ApplicationException
5
7
 
6
8
  actions_map = {
7
9
  "GET": "read",
@@ -12,7 +14,7 @@ actions_map = {
12
14
 
13
15
 
14
16
  class GatewayAdapter(Adapter):
15
- def marshal(self, result: list[dict]):
17
+ def marshal(self, result: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
16
18
  """
17
19
  Marshal the result into a event response
18
20
 
@@ -24,7 +26,7 @@ class GatewayAdapter(Adapter):
24
26
  """
25
27
  return super().marshal(result)
26
28
 
27
- def unmarshal(self, event):
29
+ def unmarshal(self, event: Dict[str, Any]) -> Operation:
28
30
  """
29
31
  Get parameters from the Lambda event.
30
32
 
@@ -34,8 +36,15 @@ class GatewayAdapter(Adapter):
34
36
  Returns:
35
37
  - tuple: Tuple containing data, query and metadata parameters.
36
38
  """
37
- entity = event.get("resource").split("/")[1]
38
- action = actions_map.get(event.get("httpMethod").upper(), "read")
39
+ resource = event.get("resource")
40
+ if resource is not None and "/" in resource:
41
+ parts = resource.split("/")
42
+ entity = parts[1] if len(parts) > 1 else None
43
+ else:
44
+ entity = None
45
+
46
+ method = str(event.get("httpMethod", "")).upper()
47
+ action = actions_map.get(method, "read")
39
48
 
40
49
  event_params = {}
41
50
 
@@ -56,11 +65,66 @@ class GatewayAdapter(Adapter):
56
65
  if body is not None and len(body) > 0:
57
66
  store_params = json.loads(body)
58
67
 
59
- roles = []
60
68
  authorizer_info = event.get("requestContext", {}).get("authorizer", {})
61
69
  claims = authorizer_info.get("claims", {})
62
- roles = claims.get("roles", [])
63
- subject = claims.get("subject")
70
+
71
+ # Decode JSON-encoded arrays from OAuth context
72
+ roles_raw = claims.get("roles", [])
73
+ if isinstance(roles_raw, str):
74
+ try:
75
+ roles = json.loads(roles_raw)
76
+ except (json.JSONDecodeError, TypeError):
77
+ roles = []
78
+ else:
79
+ roles = roles_raw if isinstance(roles_raw, list) else []
80
+
81
+ groups_raw = claims.get("groups", [])
82
+ if isinstance(groups_raw, str):
83
+ try:
84
+ groups = json.loads(groups_raw)
85
+ except (json.JSONDecodeError, TypeError):
86
+ groups = []
87
+ else:
88
+ groups = groups_raw if isinstance(groups_raw, list) else []
89
+
90
+ permissions_raw = claims.get("permissions", [])
91
+ if isinstance(permissions_raw, str):
92
+ try:
93
+ permissions = json.loads(permissions_raw)
94
+ except (json.JSONDecodeError, TypeError):
95
+ permissions = []
96
+ else:
97
+ permissions = permissions_raw if isinstance(permissions_raw, list) else []
98
+
99
+ subject = claims.get("sub")
100
+ scope_str = claims.get("scope")
101
+
102
+ # Enforce OAuth scopes (simulating API Gateway authorizer behavior)
103
+ # Required scope pattern: read|write|delete:<entity>
104
+ if entity and scope_str:
105
+ required_action = {
106
+ "GET": "read",
107
+ "POST": "write",
108
+ "PUT": "write",
109
+ "PATCH": "write",
110
+ "DELETE": "delete",
111
+ }.get(method, "read")
112
+ required_scope = f"{required_action}:{entity}"
113
+ token_scopes = set(str(scope_str).split())
114
+
115
+ def _has_scope(required: str) -> bool:
116
+ return (
117
+ required in token_scopes
118
+ or f"{required_action}:*" in token_scopes
119
+ or "*" in token_scopes
120
+ or "*:*" in token_scopes
121
+ )
122
+
123
+ if not _has_scope(required_scope):
124
+ raise ApplicationException(
125
+ 401,
126
+ ("insufficient_scope: required_scope=" + required_scope),
127
+ )
64
128
 
65
129
  return Operation(
66
130
  entity=entity,
@@ -69,11 +133,15 @@ class GatewayAdapter(Adapter):
69
133
  query_params=query_params,
70
134
  metadata_params=metadata_params,
71
135
  roles=roles,
136
+ groups=groups,
72
137
  subject=subject,
138
+ permissions=permissions,
73
139
  claims=claims,
74
140
  )
75
141
 
76
- def _convert_parameters(self, parameters):
142
+ def _convert_parameters(
143
+ self, parameters: Optional[Dict[str, Any]]
144
+ ) -> Optional[Dict[str, Any]]:
77
145
  """
78
146
  Convert parameters to appropriate types.
79
147
 
@@ -97,7 +165,9 @@ class GatewayAdapter(Adapter):
97
165
  result[parameter] = value
98
166
  return result
99
167
 
100
- def split_params(self, parameters: dict):
168
+ def split_params(
169
+ self, parameters: Dict[str, Any]
170
+ ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
101
171
  """
102
172
  Split a dictionary into two dictionaries based on keys.
103
173
 
@@ -193,20 +193,43 @@ class SQLQueryHandler:
193
193
  log.info(f"role: {role}, role_permissions: {role_permissions}")
194
194
  if len(role_permissions) == 0:
195
195
  continue
196
+
197
+ # Extract permission patterns for this role
198
+ read_perm = role_permissions.get("read", "")
199
+ write_perm = role_permissions.get("write", "")
200
+ read_pattern = self._extract_permission_pattern(read_perm)
201
+ write_pattern = self._extract_permission_pattern(write_perm)
202
+
196
203
  for prop_name, property in properties.items():
197
204
  log.info(f"prop_name: {prop_name}, property: {property}")
198
- if permission_type == "read" and re.match(
199
- role_permissions.get("read", ""), prop_name
200
- ):
205
+ if permission_type == "read" and re.match(read_pattern, prop_name):
201
206
  allowed_properties[prop_name] = property
202
- if permission_type == "write" and re.match(
203
- role_permissions.get("write", ""), prop_name
204
- ):
207
+ if permission_type == "write" and re.match(write_pattern, prop_name):
205
208
  allowed_properties[prop_name] = property
206
209
 
207
210
  log.info(f"allowed_properties: {allowed_properties}")
208
211
  return allowed_properties
209
212
 
213
+ def _extract_permission_pattern(self, permission_rule) -> str:
214
+ """Extract the regex pattern from a permission rule.
215
+
216
+ Args:
217
+ permission_rule: Can be a string (regex) or dict with
218
+ 'properties'/'fields' key
219
+
220
+ Returns:
221
+ str: The regex pattern to match property names
222
+ """
223
+ if isinstance(permission_rule, str):
224
+ return permission_rule
225
+ elif isinstance(permission_rule, dict):
226
+ # Support both 'properties' (preferred) and 'fields' (legacy)
227
+ return permission_rule.get("properties") or permission_rule.get(
228
+ "fields", ""
229
+ )
230
+ else:
231
+ return ""
232
+
210
233
  @property
211
234
  def selection_results(self) -> Dict:
212
235
  raise NotImplementedError()
@@ -33,20 +33,34 @@ class SQLSelectSchemaQueryHandler(SQLSchemaQueryHandler):
33
33
  if "default" in perms:
34
34
  provider = perms.get("default", {}) or {}
35
35
  read_map = provider.get("read", {}) or {}
36
+ role_permissions = perms.get("default", {})
36
37
  else:
37
38
  # legacy role-first -> synthesize read map
38
39
  read_map = {}
40
+ role_permissions = perms
39
41
  for role, role_perms in perms.items():
40
42
  if isinstance(role_perms, dict):
41
43
  read_map[role] = role_perms.get("read")
42
44
 
43
45
  filters = []
44
46
  for role in self.operation.roles or []:
47
+ # Check for role-level WHERE clause (hybrid approach)
48
+ role_where = None
49
+ if isinstance(role_permissions.get(role), dict):
50
+ role_where = role_permissions[role].get("where")
51
+
52
+ # Check for operation-level WHERE clause
53
+ operation_where = None
45
54
  rule = read_map.get(role)
46
55
  if isinstance(rule, dict):
47
- where = rule.get("where")
48
- if isinstance(where, str) and where.strip():
49
- filters.append(self._template_where(where))
56
+ operation_where = rule.get("where")
57
+
58
+ # Operation-level takes precedence, fallback to role-level
59
+ where_clause = operation_where if operation_where else role_where
60
+
61
+ if isinstance(where_clause, str) and where_clause.strip():
62
+ filters.append(self._template_where(where_clause))
63
+
50
64
  if not filters:
51
65
  return ""
52
66
  return "(" + ") OR (".join(filters) + ")"
@@ -1,7 +1,7 @@
1
1
  import json
2
2
  import logging
3
3
  import os
4
- from typing import Optional, Mapping
4
+ from typing import Optional, Mapping, Any
5
5
 
6
6
  from api_foundry_query_engine.utils.api_model import set_api_model
7
7
  from api_foundry_query_engine.utils.app_exception import ApplicationException
@@ -14,8 +14,8 @@ class QueryEngine:
14
14
  def __init__(self, config: Mapping[str, str]):
15
15
  self.adapter = GatewayAdapter(config)
16
16
 
17
- def handler(self, event):
18
- log.debug(f"event: {event}")
17
+ def handler(self, event) -> dict[str, Any]:
18
+ log.debug("event: %s", event)
19
19
  try:
20
20
  response = self.adapter.process_event(event)
21
21
 
@@ -27,20 +27,20 @@ class QueryEngine:
27
27
  "body": json.dumps(response),
28
28
  }
29
29
  except ApplicationException as e:
30
- log.error(f"exception: {e}", exc_info=True)
30
+ log.error("exception: %s", e, exc_info=True)
31
31
  return {
32
32
  "isBase64Encoded": False,
33
33
  "statusCode": e.status_code,
34
34
  "headers": {"Content-Type": "application/json"},
35
- "body": json.dumps({"message": f"exception: {e}"}),
35
+ "body": json.dumps({"message": "exception: %s" % e}),
36
36
  }
37
- except Exception as e:
38
- log.error(f"exception: {e}", exc_info=True)
37
+ except RuntimeError as e:
38
+ log.error("runtime error: %s", e, exc_info=True)
39
39
  return {
40
40
  "isBase64Encoded": False,
41
41
  "statusCode": 500,
42
42
  "headers": {"Content-Type": "application/json"},
43
- "body": json.dumps({"message": f"exception: {e}"}),
43
+ "body": json.dumps({"message": f"runtime error: {e}"}),
44
44
  }
45
45
 
46
46
 
@@ -49,16 +49,14 @@ query_engine: Optional[QueryEngine] = None
49
49
 
50
50
 
51
51
  def handler(event, _):
52
- global engine_config, query_engine
53
- log.info(f"engine_config: {engine_config}")
54
- if engine_config is None:
52
+ if not hasattr(handler, "engine_config"):
55
53
  log.info("Loading engine config from environment variables")
56
- engine_config = os.environ
57
- log.info(f"engine_config: {engine_config}")
54
+ handler.engine_config = os.environ
55
+ log.info(f"engine_config: {handler.engine_config}")
58
56
 
59
- if query_engine is None:
60
- set_api_model(engine_config)
57
+ if not hasattr(handler, "query_engine"):
58
+ set_api_model(handler.engine_config)
61
59
  log.info("Creating QueryEngine instance")
62
- query_engine = QueryEngine(engine_config)
60
+ handler.query_engine = QueryEngine(handler.engine_config)
63
61
 
64
- return query_engine.handler(event)
62
+ return handler.query_engine.handler(event)
@@ -1,4 +1,4 @@
1
- from typing import Any, Dict, Optional
1
+ from typing import Any, Dict, List, Optional
2
2
 
3
3
 
4
4
  class Operation:
@@ -35,7 +35,10 @@ class Operation:
35
35
  query_params: Optional[Dict[str, Any]] = None,
36
36
  store_params: Optional[Dict[str, Any]] = None,
37
37
  metadata_params: Optional[Dict[str, Any]] = None,
38
- roles: Optional[Dict[str, Any]] = None,
38
+ roles: Optional[List[str]] = None,
39
+ scope: Optional[Dict[str, Any]] = None,
40
+ groups: Optional[List[str]] = None,
41
+ permissions: Optional[List[str]] = None,
39
42
  subject: Optional[str] = None,
40
43
  claims: Optional[Dict[str, Any]] = None,
41
44
  ):
@@ -73,7 +76,10 @@ class Operation:
73
76
  self.metadata_params = metadata_params or {}
74
77
 
75
78
  # Roles defining the context in which the operation is allowed.
76
- self.roles = roles or {}
79
+ self.scope = scope or {}
80
+ self.roles = roles or []
81
+ self.groups = groups or []
82
+ self.permissions = permissions or []
77
83
 
78
84
  self.subject = subject
79
85
  # Full set of claims (when available). Used for row-level filters
@@ -49,38 +49,117 @@ class SchemaObjectProperty:
49
49
  def convert_to_db_value(self, value: str) -> Optional[Any]:
50
50
  if value is None:
51
51
  return None
52
- conversion_mapping = {
53
- "string": lambda x: x,
54
- "number": float,
55
- "float": float,
56
- "integer": int,
57
- "boolean": lambda x: x.lower() == "true",
58
- "date": lambda x: datetime.strptime(x, "%Y-%m-%d").date() if x else None,
59
- "date-time": lambda x: datetime.fromisoformat(x) if x else None,
60
- "time": lambda x: datetime.strptime(x, "%H:%M:%S").time() if x else None,
61
- }
62
- conversion_func = conversion_mapping.get(
63
- self.column_type if self.column_type is not None else "string", lambda x: x
64
- )
65
- return conversion_func(value)
52
+
53
+ column_type = self.column_type if self.column_type is not None else "string"
54
+
55
+ # Handle string types
56
+ if column_type in ["string", "varchar", "char", "text", "uuid"]:
57
+ return value
58
+
59
+ # Handle numeric types - float/double/numeric variations
60
+ elif column_type in ["number", "float", "double", "numeric", "decimal", "real"]:
61
+ return float(value)
62
+
63
+ # Handle boolean types - can map to boolean or integer columns
64
+ elif column_type == "boolean":
65
+ return value.lower() == "true"
66
+ elif (
67
+ column_type in ["int", "integer", "smallint", "bigint"]
68
+ and hasattr(self, "api_type")
69
+ and self.api_type == "boolean"
70
+ ):
71
+ # Boolean API type mapping to integer column type
72
+ return 1 if value.lower() == "true" else 0
73
+
74
+ # Handle integer types (after boolean check to avoid conflicts)
75
+ elif column_type in [
76
+ "integer",
77
+ "int",
78
+ "bigint",
79
+ "smallint",
80
+ "serial",
81
+ "bigserial",
82
+ ]:
83
+ return int(value)
84
+
85
+ # Handle date types
86
+ elif column_type == "date":
87
+ return datetime.strptime(value, "%Y-%m-%d").date() if value else None
88
+
89
+ # Handle datetime types - various column type names
90
+ elif column_type in ["date-time", "datetime", "timestamp", "timestamptz"]:
91
+ return datetime.fromisoformat(value) if value else None
92
+
93
+ # Handle time types
94
+ elif column_type in ["time", "timetz"]:
95
+ return datetime.strptime(value, "%H:%M:%S").time() if value else None
96
+
97
+ # Default to string conversion for unknown types
98
+ else:
99
+ return value
66
100
 
67
101
  def convert_to_api_value(self, value) -> Optional[Any]:
68
102
  if value is None:
69
103
  return None
70
- conversion_mapping = {
71
- "string": lambda x: x,
72
- "number": float,
73
- "float": float,
74
- "integer": int,
75
- "boolean": str,
76
- "date": lambda x: x.date().isoformat() if x else None,
77
- "date-time": lambda x: x.isoformat() if x else None,
78
- "time": lambda x: x.time().isoformat() if x else None,
79
- }
80
- conversion_func = conversion_mapping.get(
81
- self.api_type if self.api_type is not None else "string", lambda x: x
82
- )
83
- return conversion_func(value)
104
+
105
+ api_type = self.api_type if self.api_type is not None else "string"
106
+
107
+ # Handle string types (including UUID which is represented as string in API)
108
+ if api_type in ["string", "uuid"]:
109
+ return str(value) if value is not None else None
110
+
111
+ # Handle integer type
112
+ elif api_type == "integer":
113
+ return int(value) if value is not None else None
114
+
115
+ # Handle number and float types
116
+ elif api_type in ["number", "float"]:
117
+ return float(value) if value is not None else None
118
+
119
+ # Handle boolean type - convert any value to string representation
120
+ elif api_type == "boolean":
121
+ if isinstance(value, bool):
122
+ return str(value)
123
+ elif isinstance(value, int):
124
+ # Handle boolean stored as integer (0/1)
125
+ return "true" if value != 0 else "false"
126
+ elif isinstance(value, str):
127
+ return value.lower() in ["true", "1", "yes", "on"]
128
+ else:
129
+ return str(bool(value))
130
+
131
+ # Handle date type
132
+ elif api_type == "date":
133
+ if hasattr(value, "date"):
134
+ # If it's a datetime, extract the date part
135
+ return value.date().isoformat()
136
+ elif hasattr(value, "isoformat"):
137
+ # If it's already a date
138
+ return value.isoformat()
139
+ else:
140
+ return str(value)
141
+
142
+ # Handle datetime type
143
+ elif api_type == "date-time":
144
+ if hasattr(value, "isoformat"):
145
+ return value.isoformat()
146
+ else:
147
+ return str(value)
148
+
149
+ # Handle time type
150
+ elif api_type == "time":
151
+ if hasattr(value, "time"):
152
+ # If it's a datetime, extract the time part
153
+ return value.time().isoformat()
154
+ elif hasattr(value, "isoformat"):
155
+ # If it's already a time
156
+ return value.isoformat()
157
+ else:
158
+ return str(value)
159
+
160
+ # Default to string conversion for unknown types
161
+ else:
162
+ return str(value) if value is not None else None
84
163
 
85
164
 
86
165
  class SchemaObjectAssociation:
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "api-foundry-query-engine"
7
- version = "0.8.11"
7
+ version = "0.8.30"
8
8
  authors = [
9
9
  { name="Daniel Repik", email="danrepik@icloud.com" },
10
10
  ]
@@ -1,120 +0,0 @@
1
- from api_foundry_query_engine.connectors.connection import Connection, Cursor
2
- from api_foundry_query_engine.utils.logger import logger
3
- import boto3
4
- import time
5
-
6
- # Initialize the logger
7
- log = logger(__name__)
8
-
9
-
10
- class AthenaCursor(Cursor):
11
- def __init__(self, athena_client, db_config):
12
- self.athena_client = athena_client
13
- self.db_config = db_config
14
-
15
- def execute(self, sql: str, parameters: dict, result_columns: list[str]) -> list:
16
- """
17
- Execute a SQL statement on the Athena database.
18
-
19
- Parameters:
20
- - sql (str): The SQL query to execute.
21
- - parameters (dict): Parameters for the query (currently unused; substitute manually in the query).
22
- - result_columns (list[str]): Column names for the results.
23
-
24
- Returns:
25
- - list[dict]: Query results as a list of dictionaries.
26
- """
27
- # Substitute parameters into the query manually (Athena doesn't support placeholders)
28
- for key, value in parameters.items():
29
- placeholder = f":{key}"
30
- if isinstance(value, str):
31
- value = f"'{value}'" # Wrap strings in quotes
32
- sql = sql.replace(placeholder, str(value))
33
-
34
- log.info(f"Executing query: {sql}")
35
-
36
- # Start the query execution
37
- response = self.athena_client.start_query_execution(
38
- QueryString=sql,
39
- QueryExecutionContext={"Database": self.db_config["database"]},
40
- ResultConfiguration={"OutputLocation": self.db_config["output_location"]},
41
- )
42
- query_execution_id = response["QueryExecutionId"]
43
- log.info(f"Query execution started: {query_execution_id}")
44
-
45
- # Wait for the query to complete
46
- while True:
47
- status_response = self.athena_client.get_query_execution(
48
- QueryExecutionId=query_execution_id
49
- )
50
- status = status_response["QueryExecution"]["Status"]["State"]
51
-
52
- if status in ["SUCCEEDED", "FAILED", "CANCELLED"]:
53
- break
54
- time.sleep(2)
55
-
56
- if status == "FAILED":
57
- raise Exception(
58
- f"Query failed: {status_response['QueryExecution']['Status']['StateChangeReason']}"
59
- )
60
- elif status == "CANCELLED":
61
- raise Exception("Query was cancelled.")
62
-
63
- # Fetch results
64
- results_response = self.athena_client.get_query_results(
65
- QueryExecutionId=query_execution_id
66
- )
67
- rows = results_response["ResultSet"]["Rows"]
68
-
69
- # Convert rows to dictionaries
70
- result = []
71
- for row in rows[1:]: # Skip the first row (header row)
72
- record = {
73
- col: value.get("VarCharValue", None)
74
- for col, value in zip(result_columns, row["Data"])
75
- }
76
- result.append(record)
77
-
78
- return result
79
-
80
- def close(self):
81
- log.info("Closing Athena cursor (no persistent connection to close).")
82
-
83
-
84
- class AthenaConnection(Connection):
85
- def __init__(self, db_config: dict) -> None:
86
- """
87
- Initialize the AthenaConnection with the given configuration.
88
-
89
- Parameters:
90
- - db_config (dict): A dictionary containing Athena connection configuration.
91
- Required keys:
92
- - 'region': AWS region of the Athena service.
93
- - 'output_location': S3 bucket location for query results.
94
- - 'database': Athena database name.
95
- """
96
- super().__init__(db_config)
97
- self.athena_client = self.get_athena_client()
98
-
99
- def engine(self) -> str:
100
- return "athena"
101
-
102
- def cursor(self) -> Cursor:
103
- return AthenaCursor(self.athena_client, self.db_config)
104
-
105
- def commit(self):
106
- # Athena queries are read-only; no commit operation is required
107
- log.info("Athena does not support transactions. Commit is a no-op.")
108
-
109
- def close(self):
110
- # Athena uses a stateless API, so there's nothing to close
111
- log.info("Closing Athena connection (no persistent connection to close).")
112
-
113
- def get_athena_client(self):
114
- """
115
- Get a Boto3 client for Athena.
116
- """
117
- region = self.db_config.get("region")
118
- if not region:
119
- raise ValueError("Athena configuration must include 'region'.")
120
- return boto3.client("athena", region_name=region)
@@ -1,105 +0,0 @@
1
- from api_foundry_query_engine.utils.logger import logger
2
- from api_foundry_query_engine.operation import Operation
3
- from api_foundry_query_engine.services.service import ServiceAdapter
4
-
5
- log = logger(__name__)
6
-
7
-
8
- class SecurityService(ServiceAdapter):
9
- def execute(self, operation: Operation) -> list[dict]:
10
- """
11
- Execute the operation after validating security instructions.
12
-
13
- Parameters:
14
- - operation (Operation): The operation to execute.
15
-
16
- Returns:
17
- - list[dict]: The result of the operation, filtered based on security rules.
18
- """
19
- log.debug(f"Executing operation: {operation}")
20
-
21
- # Resolve the schema or path operation to retrieve the `security` attribute
22
- schema_or_path = self.schema_or_path_resolver(
23
- operation.entity, operation.action
24
- )
25
- if not schema_or_path or "security" not in schema_or_path:
26
- raise ValueError(
27
- f"No security instructions found for entity {operation.entity}"
28
- )
29
-
30
- security_rules = schema_or_path["security"]
31
-
32
- # Validate query and store parameters
33
- self._validate_query_params(operation.query_params, security_rules)
34
- self._validate_store_params(operation.store_params, security_rules)
35
-
36
- # Execute the next service
37
- result = self.next_service.execute(operation)
38
-
39
- # Filter the result based on security rules
40
- filtered_result = self._filter_result(result, security_rules)
41
- return filtered_result
42
-
43
- def _validate_query_params(self, query_params: dict, security_rules: dict):
44
- """
45
- Validate query parameters against read permissions.
46
-
47
- Parameters:
48
- - query_params (dict): Query parameters to validate.
49
- - security_rules (dict): Security rules for validation.
50
-
51
- Raises:
52
- - PermissionError: If a query parameter violates read permissions.
53
- """
54
- read_permissions = security_rules.get("read", [])
55
- invalid_params = [
56
- key
57
- for key in query_params
58
- if key not in read_permissions and read_permissions != ["*"]
59
- ]
60
- if invalid_params:
61
- raise PermissionError(f"Query parameters not permitted: {invalid_params}")
62
-
63
- def _validate_store_params(self, store_params: dict, security_rules: dict):
64
- """
65
- Validate store parameters against write permissions.
66
-
67
- Parameters:
68
- - store_params (dict): Store parameters to validate.
69
- - security_rules (dict): Security rules for validation.
70
-
71
- Raises:
72
- - PermissionError: If a store parameter violates write permissions.
73
- """
74
- write_permissions = security_rules.get("write", [])
75
- invalid_params = [
76
- key
77
- for key in store_params
78
- if key not in write_permissions and write_permissions != ["*"]
79
- ]
80
- if invalid_params:
81
- raise PermissionError(f"Store parameters not permitted: {invalid_params}")
82
-
83
- def _filter_result(self, result: list[dict], security_rules: dict) -> list[dict]:
84
- """
85
- Filter the result based on read permissions.
86
-
87
- Parameters:
88
- - result (list[dict]): The original result set.
89
- - security_rules (dict): Security rules for filtering.
90
-
91
- Returns:
92
- - list[dict]: The filtered result set.
93
- """
94
- read_permissions = security_rules.get("read", [])
95
- if read_permissions == ["*"]:
96
- return result # No filtering needed
97
-
98
- filtered_result = []
99
- for record in result:
100
- filtered_record = {
101
- key: value for key, value in record.items() if key in read_permissions
102
- }
103
- filtered_result.append(filtered_record)
104
-
105
- return filtered_result