To avoid the risk of HTTP-call's running indefinitely, optional parameter http_timeout can now be specified when creating a task. Default timeout is set to (30.05, 300) and is applied to all HTTP-calls.

Author: sdktjj

src/sas_airflow_provider/hooks/sas.py

@@ -7,7 +7,6 @@
 import urllib3
 from urllib3.exceptions import InsecureRequestWarning
 
-
 class SasHook(BaseHook):
     """Hook to manage connection to SAS"""
 
@@ -29,7 +28,7 @@ def __init__(self, conn_id: str = None) -> None:
         self.cert_verify = True
         self.grant_type = None
 
-    def get_conn(self):
+    def get_conn(self, http_timeout=None):
         """Returns a SAS connection."""
         if self.conn_id is None:
             self.conn_id = self.default_conn_name
@@ -55,11 +54,11 @@ def get_conn(self):
             self.log.info("Using custom TLS CA certificate bundle file")
 
         if not self.sas_conn:
-            self.sas_conn = self._create_session_for_connection()
+            self.sas_conn = self._create_session_for_connection(http_timeout=http_timeout)
 
         return self.sas_conn
 
-    def _create_session_for_connection(self):
+    def _create_session_for_connection(self, http_timeout=None):
         self.log.info(f"Creating session for connection named %s to host %s",
                       self.conn_id,
                       self.host)
@@ -82,12 +81,15 @@ def _create_session_for_connection(self):
 
             self.log.info("Get oauth token (see README if this crashes)")
             response = requests.post(
-                f"{self.host}/SASLogon/oauth/token", data=payload, verify=self.cert_verify, headers=my_headers
+                    f"{self.host}/SASLogon/oauth/token", 
+                    data=payload, 
+                    verify=self.cert_verify,
+                    headers=my_headers,
+                    timeout=http_timeout
             )
-
             if response.status_code != 200:
-                raise RuntimeError(f"Get token failed: {response.text}")
-
+                raise RuntimeError(f"Get token failed with status code: {response.status_code}") 
+                
             r = response.json()
             self.token = r["access_token"]
 

src/sas_airflow_provider/operators/sas_studio.py

@@ -24,7 +24,7 @@
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from sas_airflow_provider.hooks.sas import SasHook
-from sas_airflow_provider.util.util import dump_logs, stream_log, create_or_connect_to_session, end_compute_session
+from sas_airflow_provider.util.util import stream_log, create_or_connect_to_session, end_compute_session
 
 # main API URI for Code Gen
 URI_BASE = "/studioDevelopment/code"
@@ -88,6 +88,8 @@ class SASStudioOperator(BaseOperator):
         temporary unobtainable. When unknown_state_timeout is reached without the state being retrievable, the operator 
         will throw an AirflowFailException and the task will be marked as failed. 
         Default value is 0, meaning the task will fail immediately if the state could not be retrieved.
+    :para http_timeout: (optional) Timeout for https requests. Default value is (30.05, 300), meaning a connect timeout sligthly above 30 seoconds and 
+        a read timeout of 300 seconds where the operator will wait for the server to send a response.
     """
 
     ui_color = "#CCE5FF"
@@ -113,6 +115,7 @@ def __init__(
             compute_session_id="",
             output_macro_var_prefix="",
             unknown_state_timeout=0,
+            http_timeout=(30.05, 300),
             **kwargs,
     ) -> None:
 
@@ -146,6 +149,9 @@ def __init__(
         self.on_failure_callback=[on_failure]
         self.on_retry_callback=[on_retry]
 
+        # Timeout
+        self.http_timeout=http_timeout
+
 
     def execute(self, context):
         if self.path_type not in ['compute', 'content', 'raw']:
@@ -158,13 +164,14 @@ def execute(self, context):
         try:
             self.log.info("Authenticate connection")
             h = SasHook(self.connection_name)
-            self.connection = h.get_conn()
+            self.connection = h.get_conn(http_timeout=self.http_timeout)
 
             # Create compute session
             if not self.compute_session_id:
                 compute_session = create_or_connect_to_session(self.connection,
                                                             self.compute_context_name, 
-                                                            AIRFLOW_SESSION_NAME if self.allways_reuse_session else None)
+                                                            AIRFLOW_SESSION_NAME if self.allways_reuse_session else None, 
+                                                            http_timeout=self.http_timeout)
                 self.compute_session_id = compute_session["id"]
             else:
                 self.log.info(f"Compute Session {self.compute_session_id} was provided")
@@ -240,7 +247,7 @@ def _clean_up(self, also_kill_reused_session=False):
             if (also_kill_reused_session and self.allways_reuse_session) or self.allways_reuse_session==False:
                 try:
                     self.log.info(f"Deleting session with id {self.compute_session_id}")
-                    success_end = end_compute_session(self.connection, self.compute_session_id)
+                    success_end = end_compute_session(self.connection, self.compute_session_id, http_timeout=self.http_timeout)
                     if success_end:
                         self.log.info(f"Compute session succesfully deleted")
                     else:
@@ -300,7 +307,7 @@ def _generate_object_code(self):
             "wrapperCode": self.codegen_wrap_code,
         }
 
-        response = self.connection.post(uri, json=req)
+        response = self.connection.post(uri, json=req, timeout=self.http_timeout)
         if not response.ok:
             raise RuntimeError(f"Code generation failed: {response.text}")
 
@@ -311,7 +318,7 @@ def _run_job_and_wait(self, job_request: dict, poll_interval: int) -> (dict, boo
 
         #Kick off job request. if failures, no harm is done.
         try:
-            response = self.connection.post(uri, json=job_request)
+            response = self.connection.post(uri, json=job_request, timeout=self.http_timeout)
         except Exception as e:
             raise AirflowException(f"Error when creating Job Request {e}") 
 
@@ -336,7 +343,7 @@ def _run_job_and_wait(self, job_request: dict, poll_interval: int) -> (dict, boo
             time.sleep(poll_interval)
 
             try:
-                response = self.connection.get(uri)
+                response = self.connection.get(uri, timeout=self.http_timeout)
                 if not response.ok:
                     countUnknownState = countUnknownState + 1
                     self.log.info(f'Invalid response code {response.status_code} from {uri}. Will set state=unknown and continue checking...')
@@ -352,7 +359,7 @@ def _run_job_and_wait(self, job_request: dict, poll_interval: int) -> (dict, boo
 
                     # Get the latest new log lines.
                     if self.exec_log and state != "unknown":
-                        num_log_lines=stream_log(self.connection, job, num_log_lines)
+                        num_log_lines=stream_log(self.connection, job, num_log_lines, http_timeout=self.http_timeout)
 
             except Exception as e:
                 countUnknownState = countUnknownState + 1
@@ -365,7 +372,7 @@ def _run_job_and_wait(self, job_request: dict, poll_interval: int) -> (dict, boo
 
         # Be sure to Get the latest new log lines after the job have finished.
         if self.exec_log:
-            num_log_lines=stream_log(self.connection, job, num_log_lines)
+            num_log_lines=stream_log(self.connection, job, num_log_lines, http_timeout=self.http_timeout)
 
         self.log.info("Job request has completed execution with the status: " + str(state))
         success = True
@@ -381,7 +388,7 @@ def _set_output_variables(self, context):
 
         # retrieve variables from compute session
         uri = f"/compute/sessions/{self.compute_session_id}/variables?limit=999&filter=startsWith(name,'{self.output_macro_var_prefix}')"
-        response = self.connection.get(uri, headers={'Accept': '*/*'})
+        response = self.connection.get(uri, headers={'Accept': '*/*'}, timeout=self.http_timeout)
         if not response.ok:
             raise RuntimeError(f"get compute variables failed with {response.status_code}")
         v = response.json()["items"]

src/sas_airflow_provider/util/util.py

@@ -22,37 +22,38 @@
 import logging
 
 
-def get_folder_file_contents(session, path: str) -> str:
+def get_folder_file_contents(session, path: str, http_timeout=None) -> str:
     """
     Fetch a file from folder service
     :param session:
     :param path:
+    :param http_timeout: Timeout for http connection
     :return:
     """
     member = get_member_by_path(session, path)
     if member['contentType'] != 'file':
         raise RuntimeError(f"folder item is not a file: '{path}'")
 
     uri = member['uri'] + '/content'
-    response = session.get(uri)
+    response = session.get(uri, timeout=http_timeout)
     if not response.ok:
         raise RuntimeError(f"File {path} was not found or could not be accessed. error code: {response.status_code}")
 
     return response.text
 
 
-def get_folder_by_path(session, path: str) -> dict:
+def get_folder_by_path(session, path: str, http_timeout=None) -> dict:
     """
     Get a folder given the path.
     Return a folder object, or raise an error
     """
-    response = session.get('/folders/folders/@item', params={'path': path})
+    response = session.get('/folders/folders/@item', params={'path': path}, timeout=http_timeout)
     if response.ok:
         return response.json()
     raise RuntimeError(response.text)
 
 
-def get_member_by_path(session, path: str) -> dict:
+def get_member_by_path(session, path: str, http_timeout=None) -> dict:
     """
     Get a folder member given the full path.
     Return a folder member (object), or an empty dict if not found
@@ -61,12 +62,12 @@ def get_member_by_path(session, path: str) -> dict:
     if len(parts) < 2:
         raise RuntimeError(f"invalid path '{path}'")
 
-    f = get_folder_by_path(session, parts[0])
+    f = get_folder_by_path(session, parts[0], http_timeout=http_timeout)
 
     uri = get_uri(f['links'], 'members')
     if not uri:
         raise RuntimeError("failed to find members uri link")
-    response = session.get(uri, params={'filter': f'eq("name","{parts[1]}")'})
+    response = session.get(uri, params={'filter': f'eq("name","{parts[1]}")'}, timeout=http_timeout)
 
     if not response.ok:
         raise RuntimeError(f"failed to get folder members for '{path}'")
@@ -79,18 +80,19 @@ def get_member_by_path(session, path: str) -> dict:
     return member
 
 
-def get_compute_session_file_contents(session, compute_session, path: str) -> str:
+def get_compute_session_file_contents(session, compute_session, path: str, http_timeout=None) -> str:
     """
     Fetch a file from the compute session file system
     :param session: the rest session that includes auth token
     :param compute_session: the compute session id
     :param path: full path to the file in the file system
+    :param http_timeout: Timeout for http connection
     :return: contents of the file
     """
     p = f'{path.replace("/", "~fs~")}'
     uri = f'/compute/sessions/{compute_session}/files/{p}/content'
 
-    response = session.get(uri, headers={"Accept": "application/octet-stream"})
+    response = session.get(uri, headers={"Accept": "application/octet-stream"}, timeout=http_timeout)
     if response.ok:
         return response.text
     raise RuntimeError(f"File {path} was not found or could not be accessed. error code: {response.status_code}")
@@ -107,19 +109,19 @@ def get_uri(links, rel):
     return link["uri"]
 
 
-def stream_log(session,job,start,limit=99999) -> int:
+def stream_log(session,job,start,limit=99999, http_timeout=None) -> int:
     current_line=start
 
     log_uri = get_uri(job["links"], "log")
     if not log_uri:
-        logging.getLogger(name=None).warning("Warning: failed to retrieve log URI from links")
+        logging.getLogger(name=None).warning("Warning: failed to retrieve log URI. Maybe the log is too large.")
     else:
         try:
             # Note if it is a files link (it will be that when the job have finished), this does not support the 'start' parameter, so we need to filter it by ourself.
             # We will ignore the limit parameter in that case
             is_files_link=log_uri.startswith("/files/")
 
-            r = session.get(f"{log_uri}/content?start={start}&limit={limit}")
+            r = session.get(f"{log_uri}/content?start={start}&limit={limit}", timeout=http_timeout)
             if r.ok:
                 # Parse the json log format and print each line
                 log_contents = r.text            
@@ -134,26 +136,27 @@ def stream_log(session,job,start,limit=99999) -> int:
 
                     lines=lines+1
             else:
-                logging.getLogger(name=None).warning(f"Failed to retrieve part of the log from URI: {log_uri}/content ")
+                logging.getLogger(name=None).warning(f"Failed to retrieve parts of the log with status code {r.status_code} from URI: {log_uri}/content. Maybe the log is too large.")
         except Exception as e:
-            logging.getLogger(name=None).warning("Unable to retrieve parts of the log.")
+            logging.getLogger(name=None).warning(f"Unable to retrieve parts of the log: {e}. Maybe the log is too large.")
 
     return current_line
 
 
 
-def dump_logs(session, job):
+def dump_logs(session, job, http_timeout=None):
     """
     Get the log from the job object
     :param session: rest session
     :param job: job object that should contain links object
+    :param http_timeout: Timeout for http connection
     """
 
     log_uri = get_uri(job["links"], "log")
     if not log_uri:
         print("Warning: failed to retrieve log uri from links. Log will not be displayed")
     else:
-        r = session.get(f"{log_uri}/content")
+        r = session.get(f"{log_uri}/content", timeout=http_timeout)
         if not r.ok:
             print("Warning: failed to retrieve log content. Log will not be displayed")
 
@@ -165,9 +168,9 @@ def dump_logs(session, job):
             if t != "title":
                 print(f'{line["line"]}')
 
-def find_named_compute_session(session: requests.Session, name: str) -> dict:
+def find_named_compute_session(session: requests.Session, name: str, http_timeout=None) -> dict:
     # find session with given name
-    response = session.get(f"/compute/sessions?filter=eq(name, {name})")
+    response = session.get(f"/compute/sessions?filter=eq(name, {name})", timeout=http_timeout)
     if not response.ok:
         raise RuntimeError(f"Find sessions failed: {response.status_code}")
     sessions = response.json()
@@ -176,18 +179,19 @@ def find_named_compute_session(session: requests.Session, name: str) -> dict:
         return sessions["items"][0]
     return {}
 
-def create_or_connect_to_session(session: requests.Session, context_name: str, name = None) -> dict:
+def create_or_connect_to_session(session: requests.Session, context_name: str, name = None, http_timeout=None) -> dict:
     """
     Connect to an existing compute session by name. If that named session does not exist,
     one is created using the context name supplied
     :param session: rest session that includes oauth token
     :param context_name: the context name to use to create the session if the session was not found
     :param name: name of session to find
+    :param http_timeout: Timeout for http connection
     :return: session object
 
     """
     if name != None:
-        compute_session = find_named_compute_session(session, name)
+        compute_session = find_named_compute_session(session, name, http_timeout=http_timeout)
         if compute_session:
             return compute_session
 
@@ -197,7 +201,7 @@ def create_or_connect_to_session(session: requests.Session, context_name: str, n
 
 
     # find compute context
-    response = session.get("/compute/contexts", params={"filter": f'eq("name","{context_name}")'})
+    response = session.get("/compute/contexts", params={"filter": f'eq("name","{context_name}")'},timeout=http_timeout)
     if not response.ok:
         raise RuntimeError(f"Find context named {context_name} failed: {response.status_code}")
     context_resp = response.json()
@@ -216,7 +220,7 @@ def create_or_connect_to_session(session: requests.Session, context_name: str, n
     headers = {"Content-Type": "application/vnd.sas.compute.session.request+json"}
 
     req = json.dumps(session_request)
-    response = session.post(uri, data=req, headers=headers)
+    response = session.post(uri, data=req, headers=headers, timeout=http_timeout)
 
     if response.status_code != 201:
         raise RuntimeError(f"Failed to create session: {response.text}")
@@ -226,9 +230,9 @@ def create_or_connect_to_session(session: requests.Session, context_name: str, n
 
     return json_response
 
-def end_compute_session(session: requests.Session, id):
+def end_compute_session(session: requests.Session, id, http_timeout=None):
     uri = f'/compute/sessions/{id}'
-    response = session.delete(uri)
+    response = session.delete(uri, timeout=http_timeout)
     if not response.ok:
         return False
     return True