44from typing import Any , Dict , List , Optional , Union
55from dataclasses import dataclass , asdict , field
66import urllib .parse
7+ import math
8+ import time
79from ..core .dedupe import Dedupe
810from ..core .enums import unknown_enum_value
911from ..utils import IntegrationType , Utils
12+ from ..exceptions import APIFailure , APITimeout
1013
1114log = logging .getLogger ("socketdev" )
1215
@@ -960,15 +963,65 @@ def stream_diff(
960963 )
961964 return {}
962965
963- def stream (self , org_slug : str , full_scan_id : str , use_types : bool = False ) -> Union [dict , FullScanStreamResponse ]:
964- path = "orgs/" + org_slug + "/full-scans/" + full_scan_id
965- response = self .api .do_request (path = path , method = "GET" )
966+ def stream (
967+ self ,
968+ org_slug : str ,
969+ full_scan_id : str ,
970+ use_types : bool = False ,
971+ * ,
972+ poll_timeout_seconds : Optional [float ] = None ,
973+ poll_initial_interval_seconds : float = 5.0 ,
974+ poll_max_interval_seconds : float = 10.0 ,
975+ ) -> Union [dict , FullScanStreamResponse ]:
976+ request_timeout = self .api .request_timeout
977+ timeout = request_timeout if poll_timeout_seconds is None else poll_timeout_seconds
978+ values = (timeout , poll_initial_interval_seconds , poll_max_interval_seconds , request_timeout )
979+ if any (
980+ isinstance (value , bool )
981+ or not isinstance (value , (int , float ))
982+ or not math .isfinite (value )
983+ or value <= 0
984+ for value in values
985+ ):
986+ raise ValueError ("Full scan polling timeout and intervals must be positive finite numbers" )
987+ if poll_initial_interval_seconds > poll_max_interval_seconds :
988+ raise ValueError ("Initial polling interval must not exceed maximum polling interval" )
989+
990+ path = (
991+ "orgs/" + urllib .parse .quote (org_slug , safe = "" ) + "/full-scans/"
992+ + urllib .parse .quote (full_scan_id , safe = "" ) + "?cached=true"
993+ )
994+ deadline = time .monotonic () + timeout
995+ interval = poll_initial_interval_seconds
996+ while True :
997+ retry_error = None
998+ try :
999+ response = self .api .do_request (
1000+ path = path ,
1001+ method = "GET" ,
1002+ request_timeout = min (float (request_timeout ), max (deadline - time .monotonic (), 0.001 )),
1003+ )
1004+ except APIFailure as error :
1005+ if not error .is_transient_error ():
1006+ raise
1007+ retry_error = error
1008+ else :
1009+ if response .status_code != 202 :
1010+ break
1011+ response .close ()
1012+ # Give up now if the next attempt could not start before the deadline,
1013+ # rather than sleeping out the rest of the budget and never polling again.
1014+ if deadline - time .monotonic () <= interval :
1015+ raise APITimeout ("Timed out waiting for full scan results" ) from retry_error
1016+ time .sleep (interval )
1017+ interval = min (poll_max_interval_seconds , interval * 1.5 )
9661018
9671019 if response .status_code == 200 :
9681020 try :
9691021 stream_str = []
9701022 artifacts = {}
9711023 result = response .text
1024+ response .close ()
9721025 result = result .strip ('"' ).strip ()
9731026 for line in result .split ("\n " ):
9741027 if line != '"' and line != "" and line is not None :
@@ -993,6 +1046,7 @@ def stream(self, org_slug: str, full_scan_id: str, use_types: bool = False) -> U
9931046 return artifacts
9941047
9951048 except Exception as e :
1049+ response .close ()
9961050 error_message = f"Error parsing stream response: { str (e )} "
9971051 log .error (error_message )
9981052 if use_types :
@@ -1001,7 +1055,10 @@ def stream(self, org_slug: str, full_scan_id: str, use_types: bool = False) -> U
10011055 )
10021056 return {}
10031057
1004- error_message = response .json ().get ("error" , {}).get ("message" , "Unknown error" )
1058+ try :
1059+ error_message = response .json ().get ("error" , {}).get ("message" , "Unknown error" )
1060+ finally :
1061+ response .close ()
10051062 log .error (f"Error streaming full scan: { response .status_code } , message: { error_message } " )
10061063 if use_types :
10071064 return FullScanStreamResponse .from_dict (
0 commit comments