Skip to content

Commit 07fb42f

Browse files
committed
Apply black 26.3.1 auto-formatting
Reformat 40 files to match the new style enforced by black 26.3.1 (version bumped in the previous commit to fix CVE-2026-32274). Also align the pre-commit hook rev to 26.3.1 so local hooks and CI use the same formatter version.
1 parent 61ffc64 commit 07fb42f

41 files changed

Lines changed: 556 additions & 539 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.pre-commit-config.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
repos:
22
- repo: https://github.com/ambv/black
3-
rev: 23.1.0
3+
rev: 26.3.1
44
hooks:
55
- id: black
66
language_version: python3

enforcer/dal/robusta_config.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,10 @@ class RobustaConfig(BaseModel):
66
sinks_config: List[Dict[str, Dict]]
77
global_config: dict
88

9+
910
class RobustaToken(BaseModel):
1011
store_url: str
1112
api_key: str
1213
account_id: str
1314
email: str
14-
password: str
15+
password: str

enforcer/dal/supabase_dal.py

Lines changed: 12 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,7 @@ def __init__(self):
4545
if not self.enabled:
4646
logging.info("Not connecting to Robusta platform - robusta token not provided")
4747
return
48-
logging.info(
49-
f"Initializing Robusta platform connection for account {self.account_id} cluster {self.cluster}"
50-
)
48+
logging.info(f"Initializing Robusta platform connection for account {self.account_id} cluster {self.cluster}")
5149
options = ClientOptions(postgrest_client_timeout=SUPABASE_TIMEOUT_SECONDS)
5250
self.client = create_client(self.url, self.api_key, options)
5351
self.user_id = self.sign_in()
@@ -67,9 +65,7 @@ def execute_with_retry(_self):
6765
message = exc.message or ""
6866
if exc.code == "PGRST301" or "expired" in message.lower():
6967
# JWT expired. Sign in again and retry the query
70-
logging.error(
71-
"JWT token expired/invalid, signing in to Supabase again"
72-
)
68+
logging.error("JWT token expired/invalid, signing in to Supabase again")
7369
self.sign_in()
7470
# update the session to the new one, after re-sign in
7571
_self.session = self.client.postgrest.session
@@ -81,7 +77,7 @@ def execute_with_retry(_self):
8177
SyncQueryRequestBuilder.execute = execute_with_retry
8278

8379
@staticmethod
84-
def __load_robusta_config() -> (Optional[RobustaToken],Optional[str]):
80+
def __load_robusta_config() -> (Optional[RobustaToken], Optional[str]):
8581
config_file_path = ROBUSTA_CONFIG_PATH
8682
env_ui_token = os.environ.get("ROBUSTA_UI_TOKEN")
8783
cluster_name = os.environ.get("CLUSTER_NAME")
@@ -92,9 +88,7 @@ def __load_robusta_config() -> (Optional[RobustaToken],Optional[str]):
9288
decoded = base64.b64decode(env_ui_token)
9389
return RobustaToken(**json.loads(decoded)), cluster_name
9490
except binascii.Error:
95-
raise Exception(
96-
"binascii.Error encountered. The Robusta UI token is not a valid base64."
97-
)
91+
raise Exception("binascii.Error encountered. The Robusta UI token is not a valid base64.")
9892
except json.JSONDecodeError:
9993
raise Exception(
10094
"json.JSONDecodeError encountered. The Robusta UI token could not be parsed as JSON after being base64 decoded."
@@ -112,10 +106,7 @@ def __load_robusta_config() -> (Optional[RobustaToken],Optional[str]):
112106
if "robusta_sink" in conf.keys():
113107
token = conf["robusta_sink"].get("token")
114108
if not token:
115-
raise Exception(
116-
"No robusta token provided.\n"
117-
"Please set a valid Robusta UI token.\n "
118-
)
109+
raise Exception("No robusta token provided.\n" "Please set a valid Robusta UI token.\n ")
119110
env_replacement_token = get_env_replacement(token)
120111
if env_replacement_token:
121112
token = env_replacement_token
@@ -131,9 +122,7 @@ def __load_robusta_config() -> (Optional[RobustaToken],Optional[str]):
131122
decoded = base64.b64decode(token)
132123
return RobustaToken(**json.loads(decoded)), config.global_config.get("cluster_name")
133124
except binascii.Error:
134-
raise Exception(
135-
"binascii.Error encountered. The robusta token provided is not a valid base64."
136-
)
125+
raise Exception("binascii.Error encountered. The robusta token provided is not a valid base64.")
137126
except json.JSONDecodeError:
138127
raise Exception(
139128
"json.JSONDecodeError encountered. The Robusta token provided could not be parsed as JSON after being base64 decoded."
@@ -167,12 +156,8 @@ def __init_config(self) -> bool:
167156

168157
def sign_in(self) -> str:
169158
logging.info("Supabase DAL login")
170-
res = self.client.auth.sign_in_with_password(
171-
{"email": self.email, "password": self.password}
172-
)
173-
self.client.auth.set_session(
174-
res.session.access_token, res.session.refresh_token
175-
)
159+
res = self.client.auth.sign_in_with_password({"email": self.email, "password": self.password})
160+
self.client.auth.set_session(res.session.access_token, res.session.refresh_token)
176161
self.client.postgrest.auth(res.session.access_token)
177162
return res.user.id
178163

@@ -200,7 +185,7 @@ def get_latest_krr_scan(self, current_scan_id: Optional[str]) -> (Optional[str],
200185
latest_scan_data = sorted_scans[0]
201186
else:
202187
latest_scan_data = scans_meta_response.data[0]
203-
188+
204189
latest_scan_id = latest_scan_data["scan_id"]
205190

206191
if latest_scan_id == current_scan_id:
@@ -211,7 +196,9 @@ def get_latest_krr_scan(self, current_scan_id: Optional[str]) -> (Optional[str],
211196
scan_datetime = datetime.fromisoformat(scan_start)
212197
max_age = timedelta(hours=SCAN_AGE_HOURS_THRESHOLD)
213198
if datetime.now(timezone.utc) - scan_datetime > max_age:
214-
logging.warning(f"Latest scan {latest_scan_id} is too old (started {scan_start}). No fresh KRR scan available.")
199+
logging.warning(
200+
f"Latest scan {latest_scan_id} is too old (started {scan_start}). No fresh KRR scan available."
201+
)
215202
return None, None
216203

217204
scans_results_response = (
@@ -229,6 +216,3 @@ def get_latest_krr_scan(self, current_scan_id: Optional[str]) -> (Optional[str],
229216
except Exception:
230217
logging.exception("Supabase error while retrieving krr scan data")
231218
return None, None
232-
233-
234-

enforcer/enforcer_main.py

Lines changed: 54 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -35,44 +35,43 @@
3535
# Configure logging
3636
logger = logging.getLogger()
3737
logHandler = logging.StreamHandler(sys.stdout)
38-
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
38+
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
3939
logHandler.setFormatter(formatter)
4040
logger.addHandler(logHandler)
4141
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))
4242

4343
# Define the mention pattern regex
44-
MENTION_PATTERN = re.compile(r'@[\w.-]+')
44+
MENTION_PATTERN = re.compile(r"@[\w.-]+")
4545
ENFORCE = "enforce"
4646
IGNORE = "ignore"
4747

4848
app = FastAPI(
4949
title="KRR Enforcer mutation webhook",
5050
description="A KRR recommendations mutating webhook server for Kubernetes",
51-
version="1.0.0"
51+
version="1.0.0",
5252
)
5353

5454
dal = SupabaseDal()
5555
recommendation_store = RecommendationStore(dal)
5656
owner_store = OwnerStore()
5757

58+
5859
class AdmissionReview(BaseModel):
5960
apiVersion: str
6061
kind: str
6162
request: Dict[str, Any]
6263

64+
6365
def admission_allowed(request: AdmissionReview) -> Dict[str, Any]:
64-
return \
65-
{
66-
"apiVersion": "admission.k8s.io/v1",
67-
"kind": "AdmissionReview",
68-
"response": {
69-
"uid": request.request.get('uid'),
70-
"allowed": True
71-
}
66+
return {
67+
"apiVersion": "admission.k8s.io/v1",
68+
"kind": "AdmissionReview",
69+
"response": {"uid": request.request.get("uid"), "allowed": True},
7270
}
7371

72+
7473
def enforce_pod(pod: Dict[str, Any]) -> bool:
75-
mode = pod.get('metadata', {}).get('annotations', {}).get("admission.robusta.dev/krr-mutation-mode", None)
74+
mode = pod.get("metadata", {}).get("annotations", {}).get("admission.robusta.dev/krr-mutation-mode", None)
7675
if mode == ENFORCE:
7776
return True
7877
elif mode == IGNORE:
@@ -85,30 +84,29 @@ def enforce_pod(pod: Dict[str, Any]) -> bool:
8584
async def mutate(request: AdmissionReview):
8685
"""
8786
Handle mutating webhook requests from Kubernetes.
88-
87+
8988
Args:
9089
request (AdmissionReview): The admission review request from Kubernetes
91-
90+
9291
Returns:
9392
dict: Admission review response
9493
"""
9594
start_time = time.time()
9695
try:
9796
logging.debug("Admission request received %s", request)
9897
# Extract the object being reviewed
99-
object_to_review = request.request.get('object', {})
100-
kind = request.request.get('kind', {}).get('kind')
98+
object_to_review = request.request.get("object", {})
99+
kind = request.request.get("kind", {}).get("kind")
101100

102101
if kind == "ReplicaSet": # use create/delete admission requests, to track new/removed replica sets owners
103102
owner_store.handle_rs_admission(request.request)
104-
operation = request.request.get('operation', 'UNKNOWN')
103+
operation = request.request.get("operation", "UNKNOWN")
105104
replicaset_admissions.labels(operation=operation).inc()
106-
admission_duration.labels(kind='ReplicaSet').observe(time.time() - start_time)
105+
admission_duration.labels(kind="ReplicaSet").observe(time.time() - start_time)
107106
# Update rs_owners size metric
108107
rs_owners_size.set(owner_store.get_rs_owners_count())
109108
return admission_allowed(request)
110109

111-
112110
if kind != "Pod":
113111
logger.warning(f"Received unexpected resource mutation: {kind}")
114112
return admission_allowed(request)
@@ -144,12 +142,12 @@ async def mutate(request: AdmissionReview):
144142
logger.debug("Pod Recommendations %s", recommendations)
145143

146144
patches = []
147-
145+
148146
containers = object_to_review.get("spec", {}).get("containers", [])
149147
for i, container in enumerate(containers):
150148
container_name = container.get("name")
151149
patches.extend(patch_container_resources(i, container, recommendations.get(container_name)))
152-
150+
153151
# Record metrics for Pod mutation
154152
was_mutated = len(patches) > 0
155153
reason = "success" if was_mutated else "no_changes_needed"
@@ -166,91 +164,96 @@ async def mutate(request: AdmissionReview):
166164
response["patchType"] = "JSONPatch"
167165
response["patch"] = base64.b64encode(json.dumps(patches).encode()).decode()
168166

169-
return {
170-
"apiVersion": "admission.k8s.io/v1",
171-
"kind": "AdmissionReview",
172-
"response": response
173-
}
174-
167+
return {"apiVersion": "admission.k8s.io/v1", "kind": "AdmissionReview", "response": response}
168+
175169
except Exception as e:
176170
logger.exception("Error processing webhook request")
177171
# Record failure metric for Pod requests
178-
if request.request.get('kind', {}).get('kind') == "Pod":
172+
if request.request.get("kind", {}).get("kind") == "Pod":
179173
pod_admission_mutations.labels(mutated="false", reason="processing_error").inc()
180174
admission_duration.labels(kind="Pod").observe(time.time() - start_time)
181175
raise HTTPException(status_code=500, detail=str(e))
182176

177+
183178
@app.get("/health")
184179
async def health_check():
185180
"""
186181
Health check endpoint.
187-
182+
188183
Returns:
189184
dict: Health status
190185
"""
191186
owner_store.finalize_owner_initialization() # Init loading owners from api server, after accepting api requests
192187
return {"status": "healthy"}
193188

189+
194190
@app.get("/recommendations/{namespace}/{kind}/{name}")
195191
async def get_recommendations(namespace: str, kind: str, name: str):
196192
"""
197193
Get recommendations for a workload.
198-
194+
199195
Args:
200196
namespace: Kubernetes namespace
201197
kind: Workload kind (e.g., Deployment, StatefulSet)
202198
name: Workload name
203-
199+
204200
Returns:
205201
dict: Recommendations per container or 404 if not found
206202
"""
207203
try:
208204
recommendations: WorkloadRecommendation = recommendation_store.get_recommendations(
209205
name=name, namespace=namespace, kind=kind
210206
)
211-
207+
212208
if not recommendations:
213209
raise HTTPException(status_code=404, detail="No recommendations found for this workload")
214-
210+
215211
result = {}
216212
for container_name, container_recommendation in recommendations.container_recommendations.items():
217213
result[container_name] = {
218-
"cpu": {
219-
"request": container_recommendation.cpu.request,
220-
"limit": container_recommendation.cpu.limit
221-
} if container_recommendation.cpu else None,
222-
"memory": {
223-
"request": container_recommendation.memory.request,
224-
"limit": container_recommendation.memory.limit
225-
} if container_recommendation.memory else None
214+
"cpu": (
215+
{"request": container_recommendation.cpu.request, "limit": container_recommendation.cpu.limit}
216+
if container_recommendation.cpu
217+
else None
218+
),
219+
"memory": (
220+
{"request": container_recommendation.memory.request, "limit": container_recommendation.memory.limit}
221+
if container_recommendation.memory
222+
else None
223+
),
226224
}
227-
228-
return {
229-
"namespace": namespace,
230-
"kind": kind,
231-
"name": name,
232-
"containers": result
233-
}
234-
225+
226+
return {"namespace": namespace, "kind": kind, "name": name, "containers": result}
227+
235228
except HTTPException:
236229
raise
237230
except Exception as e:
238231
logger.exception("Error retrieving recommendations")
239232
raise HTTPException(status_code=500, detail=str(e))
240233

234+
241235
@app.get("/metrics")
242236
async def metrics():
243237
"""
244238
Prometheus metrics endpoint.
245-
239+
246240
Returns:
247241
Response: Prometheus metrics in text format
248242
"""
249243
# Update rs_owners size metric before returning metrics
250244
rs_owners_size.set(owner_store.get_rs_owners_count())
251245
return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)
252246

247+
253248
if __name__ == "__main__":
254249
import uvicorn
250+
255251
logger.info("Starting Kubernetes Webhook server on 8443...")
256-
uvicorn.run(app, host="0.0.0.0", port=8443, ssl_keyfile=ENFORCER_SSL_KEY_FILE, ssl_certfile=ENFORCER_SSL_CERT_FILE, log_level="warning")
252+
uvicorn.run(
253+
app,
254+
host="0.0.0.0",
255+
port=8443,
256+
ssl_keyfile=ENFORCER_SSL_KEY_FILE,
257+
ssl_certfile=ENFORCER_SSL_CERT_FILE,
258+
log_level="warning",
259+
)

enforcer/env_vars.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
import os
22

3-
ROBUSTA_CONFIG_PATH = os.environ.get(
4-
"ROBUSTA_CONFIG_PATH", "/etc/robusta/config/active_playbooks.yaml"
5-
)
3+
ROBUSTA_CONFIG_PATH = os.environ.get("ROBUSTA_CONFIG_PATH", "/etc/robusta/config/active_playbooks.yaml")
64
ROBUSTA_ACCOUNT_ID = os.environ.get("ROBUSTA_ACCOUNT_ID", "")
75
STORE_URL = os.environ.get("STORE_URL", "")
86
STORE_API_KEY = os.environ.get("STORE_API_KEY", "")
@@ -18,10 +16,13 @@
1816
KRR_MUTATION_MODE_DEFAULT = os.environ.get("KRR_MUTATION_MODE_DEFAULT", "enforce")
1917
REPLICA_SET_CLEANUP_INTERVAL = int(os.environ.get("REPLICA_SET_CLEANUP_INTERVAL", 600))
2018
REPLICA_SET_DELETION_WAIT = int(os.environ.get("REPLICA_SET_DELETION_WAIT", 600))
21-
SCAN_AGE_HOURS_THRESHOLD = int(os.environ.get("SCAN_AGE_HOURS_THRESHOLD", 360)) # 15 days
19+
SCAN_AGE_HOURS_THRESHOLD = int(os.environ.get("SCAN_AGE_HOURS_THRESHOLD", 360)) # 15 days
2220

2321
ENFORCER_SSL_KEY_FILE = os.environ.get("ENFORCER_SSL_KEY_FILE", "")
2422
ENFORCER_SSL_CERT_FILE = os.environ.get("ENFORCER_SSL_CERT_FILE", "")
2523

26-
EXCLUDED_CONTAINERS = [container_name.strip() for container_name
27-
in os.environ.get("EXCLUDED_CONTAINERS", "").split(",") if container_name.strip()]
24+
EXCLUDED_CONTAINERS = [
25+
container_name.strip()
26+
for container_name in os.environ.get("EXCLUDED_CONTAINERS", "").split(",")
27+
if container_name.strip()
28+
]

0 commit comments

Comments
 (0)