psr-factory 4.1.0b11__py3-none-win_amd64.whl → 5.0.0b1__py3-none-win_amd64.whl

psr/cloud/cloud.py

@@ -4,10 +4,12 @@
 
 import copy
 import functools
+import gzip
 import hashlib
 import logging
 import os
 import re
+import shutil
 import subprocess
 import sys
 import warnings
@@ -21,6 +23,7 @@ import pefile
 import zeep
 from filelock import FileLock
 
+from .aws import download_case_from_s3, upload_case_to_s3
 from .data import Case, CloudError, CloudInputError
 from .desktop import import_case
 from .log import enable_log_timestamp, get_logger
@@ -49,8 +52,12 @@ def thread_safe():
     return decorator
 
 
-def _md5sum(value: str) -> str:
-    return hashlib.md5(value.encode("utf-8")).hexdigest()  # nosec
+def _md5sum(value: str, enconding=True) -> str:
+    if enconding:
+        return hashlib.md5(value.encode("utf-8")).hexdigest()  # nosec
+    else:
+        # hash binary data
+        return hashlib.md5(value).hexdigest()  # nosec
 
 
 def hash_password(password: str) -> str:
@@ -134,6 +141,7 @@ class Client:
         self._debug_mode = kwargs.get("debug", False)
         self._dry_run = kwargs.get("dry_run", False)
         self._timeout = kwargs.get("timeout", None)
+        self._python_client = kwargs.get("python_client", False)
 
         # Client version
         self.application_version = kwargs.get("application_version", None)
@@ -151,6 +159,11 @@ class Client:
 
         self._logger.info(f"Client uid {log_id} initialized.")
 
+        if self._python_client:
+            self._logger.info(
+                "Using Python client for PSR Cloud. Some features may not be available."
+            )
+
         self._console_path_setup(**kwargs)
 
         self._credentials_setup(**kwargs)
@@ -617,20 +630,23 @@ class Client:
         if self._dry_run:
             return xml_content
 
-        self._run_console(xml_content)
-        xml = ET.parse(
-            f"{self._get_console_parent_path()}\\fake{case.program}_async.xml"
-        )
-        _check_for_errors(xml, self._logger)
-        id_parameter = xml.find("./Parametro[@nome='repositorioId']")
-        if id_parameter is None:
-            xml_str = _xml_to_str(xml)
-            raise CloudError(
-                f"Case id not found on returned XML response.\n"
-                f"Contact PSR support at psrcloud@psr-inc.com with following data:\n\n{xml_str}\n\n"
+        if self._python_client:
+            case_id = self._execute_case(parameters)
+        else:
+            self._run_console(xml_content)
+            xml = ET.parse(
+                f"{self._get_console_parent_path()}\\fake{case.program}_async.xml"
             )
+            _check_for_errors(xml, self._logger)
+            id_parameter = xml.find("./Parametro[@nome='repositorioId']")
+            if id_parameter is None:
+                xml_str = _xml_to_str(xml)
+                raise CloudError(
+                    f"Case id not found on returned XML response.\n"
+                    f"Contact PSR support at psrcloud@psr-inc.com with following data:\n\n{xml_str}\n\n"
+                )
 
-        case_id = int(id_parameter.text)
+            case_id = int(id_parameter.text)
         if not wait:
             self._logger.info(f"Case {case.name} started with id {case_id}")
 
@@ -648,7 +664,6 @@ class Client:
         return case_id
 
     def get_status(self, case_id: int) -> tuple["ExecutionStatus", str]:
-        # case = self.get_case(case_id)
         delete_xml = not self._debug_mode
         xml_content = ""
         with CreateTempFile(
@@ -666,10 +681,14 @@ class Client:
                 "comando": "obterstatusresultados",
                 "arquivoSaida": status_xml_path,
             }
-            run_xml_content = create_case_xml(parameters)
 
-            self._run_console(run_xml_content)
-            xml = ET.parse(status_xml_path)
+            xml = None
+            if self._python_client:
+                xml = self._get_status_python(parameters)
+            else:
+                run_xml_content = create_case_xml(parameters)
+                self._run_console(run_xml_content)
+                xml = ET.parse(status_xml_path)
             parameter_status = xml.find("./Parametro[@nome='statusExecucao']")
             if parameter_status is None:
                 xml_str = _xml_to_str(xml)
@@ -689,7 +708,7 @@ class Client:
             self._logger.info(f"Status: {STATUS_MAP_TEXT[status]}")
             return status, STATUS_MAP_TEXT[status]
 
-    def list_download_files(self, case_id: int) -> List[str]:
+    def list_download_files(self, case_id: int) -> List[dict]:
         xml_files = self._make_soap_request(
             "prepararListaArquivosRemotaDownload",
             "listaArquivoRemota",
@@ -719,7 +738,7 @@ class Client:
         files: Optional[List[str]] = None,
         extensions: Optional[List[str]] = None,
     ) -> None:
-        filter = ".*("
+        filter = ""
 
         if not extensions and not files:
             extensions = ["csv", "log", "hdr", "bin", "out", "ok"]
@@ -728,13 +747,12 @@ class Client:
 
         if extensions:
             Client._validate_extensions(extensions)
-            filter_elements.extend([f".*.{ext}$" for ext in extensions])
+            filter_elements.extend([f".*.{ext}" for ext in extensions])
 
         if files:
             filter_elements.extend(files)
 
         filter += "|".join(filter_elements)
-        filter += ")$"
 
         self._logger.info("Download filter: " + filter)
         case = self.get_case(case_id)
@@ -752,9 +770,15 @@ class Client:
             "filtroDownloadPorMascara": filter,
         }
 
-        xml_content = create_case_xml(parameters)
-        self._run_console(xml_content)
-        self._logger.info(f"Results downloaded to {output_path}")
+        os.makedirs(output_path, exist_ok=True)
+
+        if self._python_client:
+            self._download_results_python(parameters)  ## Not implemented yet
+        else:
+            # Download results using Console
+            xml_content = create_case_xml(parameters)
+            self._run_console(xml_content)
+            self._logger.info(f"Results downloaded to {output_path}")
 
     def cancel_case(self, case_id: int, wait: bool = False) -> bool:
         parameters = {
@@ -768,8 +792,12 @@ class Client:
             "idFila": str(case_id),
         }
 
-        xml_content = create_case_xml(parameters)
-        self._run_console(xml_content)
+        if self._python_client:
+            self._cancel_case_python(case_id, parameters)
+        else:
+            # Cancel case using Console
+            xml_content = create_case_xml(parameters)
+            self._run_console(xml_content)
         self._logger.info(f"Request to cancel case {case_id} was sent")
 
         if wait:
@@ -898,7 +926,6 @@ class Client:
         if additional_arguments:
             parameters.update(additional_arguments)
 
-        # FIXME make additional arguments work as a dictionary to work with this code
         xml_input = create_case_xml(parameters)
 
         try:
@@ -941,6 +968,284 @@ class Client:
         )
         return self._cloud_clusters_xml_cache
 
+    def _execute_case(self, case_dict) -> int:
+        """
+        Execute a case on the PSR Cloud.
+        :param case_dict: Dictionary containing the case parameters.
+        :return: Case ID of the executed case.
+        """
+        case_dict["programa"] = case_dict["modelo"]
+        case_dict["numeroProcessos"] = case_dict["nproc"]
+        case_dict["versao_cliente"] = "5.4.0"
+        filter_request_result = self._make_soap_request(
+            "obterFiltros", additional_arguments=case_dict
+        )
+        upload_filter = filter_request_result.find(
+            "./Parametro[@nome='filtroUpload']"
+        ).text
+        upload_filter = "^[a-zA-Z0-9./_]*(" + upload_filter + ")$"
+
+        # Create Repository
+        self._logger.info("Creating remote repository")
+        repository_request_result = self._make_soap_request(
+            "criarRepositorio", additional_arguments=case_dict
+        )
+
+        # Add all parameters from the XML response to case_dict
+        # Iterates over each <Parametro> element in the XML response,
+        # extracts the 'nome' attribute for the key and the element's text for the value,
+        # then adds this key-value pair to case_dict.
+        for parametro_element in repository_request_result.findall("./Parametro"):
+            param_name = parametro_element.get("nome")
+            param_value = (
+                parametro_element.text
+            )  # This will be None if the element has no text.
+            if param_name:  # Ensure the parameter has a name before adding.
+                case_dict[param_name] = param_value
+
+        repository_id = repository_request_result.find(
+            "./Parametro[@nome='repositorioId']"
+        )
+        cloud_access = repository_request_result.find(
+            "./Parametro[@nome='cloudAccess']"
+        )
+        cloud_secret = repository_request_result.find(
+            "./Parametro[@nome='cloudSecret']"
+        )
+        cloud_session_token = repository_request_result.find(
+            "./Parametro[@nome='cloudSessionToken']"
+        )
+        cloud_aws_url = repository_request_result.find("./Parametro[@nome='cloudUrl']")
+        bucket_name = repository_request_result.find(
+            "./Parametro[@nome='diretorioBase']"
+        )
+
+        self._logger.info(f"Remote repository created with ID {repository_id.text}")
+        case_dict["repositorioId"] = repository_id.text
+
+        # Filtering files to upload
+        self._logger.info("Checking list of files to send")
+
+        file_list = _filter_upload_files(case_dict["diretorioDados"], upload_filter)
+
+        if not file_list:
+            self._logger.warning(
+                "No files found to upload. Please check the upload filter."
+            )
+            return
+
+        # generating .metadata folder with checksum for each file
+        checksum_dictionary = {}
+        metadata_folder = Path(case_dict["diretorioDados"]) / ".metadata"
+        metadata_folder.mkdir(parents=True, exist_ok=True)
+        for file_path in file_list:
+            file_path = Path(file_path)
+            if not file_path.is_absolute():
+                file_path = Path(case_dict["diretorioDados"]) / file_path
+            if not file_path.exists():
+                self._logger.warning(f"File {file_path} does not exist. Skipping.")
+                continue
+            with open(file_path, "rb") as f:
+                checksum = _md5sum(f.read(), enconding=False).upper()
+            checksum_dictionary[file_path.name] = checksum
+            metadata_file = metadata_folder / (file_path.name)
+            with open(metadata_file, "w") as f:
+                f.write(checksum)
+
+        self._logger.info(
+            f"Uploading list of files to remote repository {repository_id.text}"
+        )
+
+        # Uploading files to S3
+        upload_case_to_s3(
+            files=file_list,
+            repository_id=repository_id.text,
+            cluster_name=self.cluster["name"],
+            checksums=checksum_dictionary,
+            access=cloud_access.text if cloud_access is not None else None,
+            secret=cloud_secret.text if cloud_secret is not None else None,
+            session_token=cloud_session_token.text
+            if cloud_session_token is not None
+            else None,
+            bucket_name=bucket_name.text if bucket_name is not None else None,
+            url=cloud_aws_url.text if cloud_aws_url is not None else None,
+            zip_compress=True,
+        )
+
+        self._make_soap_request(
+            "finalizarUpload",
+            additional_arguments={"repositorioId": repository_id.text},
+        )
+        self._logger.info("Files uploaded successfully. Enqueuing case.")
+        self._make_soap_request("enfileirarProcesso", additional_arguments=case_dict)
+
+        return repository_id.text
+
+    def _get_status_python(self, case_dict: dict) -> ET.Element:
+        """
+        Get the status of a case using the Python client.
+        :param case_dict: Dictionary containing the case parameters.
+        :return: XML Element with the status information.
+        """
+        try:
+            response = self._make_soap_request(
+                "obterStatusResultados", additional_arguments=case_dict
+            )
+
+            # change response "status" parameter to "statusExecucao", as it is with current PSR Cloud
+            status_param = response.find("./Parametro[@nome='status']")
+            if status_param is not None:
+                status_param.attrib["nome"] = "statusExecucao"
+
+            result_log = response.find("./Parametro[@nome='resultado']")
+            if self._verbose and result_log is not None:
+                self._logger.info(result_log.text)
+            return response
+        except Exception as e:
+            self._logger.error(f"Error getting status: {str(e)}")
+            raise CloudError(f"Failed to get status: {str(e)}")
+
+    def _cancel_case_python(self, case_id: int, xml_content: str) -> None:
+        """
+        Cancel a case using the Python client.
+        :param case_id: The ID of the case to cancel.
+        :param xml_content: XML content for the cancel request.
+        """
+        try:
+            self._make_soap_request(
+                "cancelarFila", additional_arguments={"idFila": str(case_id)}
+            )
+        except Exception as e:
+            self._logger.error(f"Error cancelling case: {str(e)}")
+            raise CloudError(f"Failed to cancel case: {str(e)}")
+
+    def _download_results_python(self, parameters: dict) -> None:
+        """
+        Download results using the Python client.
+        :param parameters: Dictionary containing the download parameters.
+        """
+
+        repository_id = parameters.get("repositorioId")
+        download_filter = parameters.get("filtroDownloadPorMascara")
+        output_path = parameters.get("diretorioDestino")
+
+        download_filter = (
+            "^[a-zA-Z0-9./_]*(" + download_filter + ")$" if download_filter else None
+        )
+        self._logger.info("Obtaining credentials for download")
+        credentials = self._make_soap_request(
+            "buscaCredenciasDownload", additional_arguments=parameters
+        )
+
+        access = credentials.find("./Parametro[@nome='cloudAccess']").text
+        secret = credentials.find("./Parametro[@nome='cloudSecret']").text
+        session_token = credentials.find("./Parametro[@nome='cloudSessionToken']").text
+        url = credentials.find("./Parametro[@nome='cloudUrl']").text
+        bucket_name = credentials.find("./Parametro[@nome='diretorioBase']").text
+        bucket_name = bucket_name.replace("repository", "repository-download")
+
+        if access is None or secret is None or session_token is None or url is None:
+            raise CloudError("Failed to retrieve credentials for downloading results.")
+
+        file_list = self.list_download_files(repository_id)
+
+        # filtering files to download
+        if download_filter:
+            filtered_file_list = [
+                file["name"]
+                for file in file_list
+                if re.match(download_filter, file["name"])
+            ]
+        else:
+            filtered_file_list = [file["name"] for file in file_list]
+
+        self._logger.info("Downloading results")
+        downloaded_list = download_case_from_s3(
+            repository_id=parameters["repositorioId"],
+            cluster_name=self.cluster["name"],
+            access=access,
+            secret=secret,
+            session_token=session_token,
+            bucket_name=bucket_name,
+            url=url,
+            output_path=output_path,
+            file_list=filtered_file_list,
+        )
+
+        # Decompress gzipped files
+        for file in filtered_file_list:
+            if self._is_file_gzipped(os.path.join(output_path, file)):
+                self._decompress_gzipped_file(os.path.join(output_path, file))
+
+        # Check if all requested files were downloaded
+        for file in filtered_file_list:
+            if file not in downloaded_list:
+                self._logger.warning(f"File {file} was not downloaded.")
+
+        self._logger.info(f"Results downloaded to {output_path}")
+
+    def _is_file_gzipped(self, file_path: str) -> bool:
+        """
+        Checks if a file is gzipped by inspecting its magic number.
+
+        :param file_path: The path to the file.
+        :return: True if the file is gzipped, False otherwise.
+        """
+        try:
+            with open(file_path, "rb") as f_check:
+                return f_check.read(2) == b"\x1f\x8b"
+        except IOError:
+            # TODO: Replace print with proper logging
+            print(
+                f"WARNING: Could not read {file_path} to check for gzip magic number."
+            )
+            return False
+
+    def _decompress_gzipped_file(self, gzipped_file_path: str) -> str:
+        """
+        Decompresses a gzipped file.
+
+        If the original filename ends with .gz, the .gz is removed for the
+        decompressed filename. Otherwise, the file is decompressed in-place.
+        The original gzipped file is removed upon successful decompression.
+
+        :param gzipped_file_path: The path to the gzipped file.
+        :return: The path to the decompressed file. If decompression fails,
+                the original gzipped_file_path is returned.
+        """
+        decompressed_target_path = (
+            gzipped_file_path[:-3]
+            if gzipped_file_path.lower().endswith(".gz")
+            else gzipped_file_path
+        )
+        # Use a temporary file for decompression to avoid data loss if issues occur
+        temp_decompressed_path = decompressed_target_path + ".decompressing_tmp"
+
+        try:
+            with gzip.open(gzipped_file_path, "rb") as f_in, open(
+                temp_decompressed_path, "wb"
+            ) as f_out:
+                shutil.copyfileobj(f_in, f_out)
+            os.remove(gzipped_file_path)
+            os.rename(temp_decompressed_path, decompressed_target_path)
+            return decompressed_target_path
+        except (gzip.BadGzipFile, EOFError, IOError) as e:
+            print(
+                f"ERROR: Failed to decompress {gzipped_file_path}: {e}. Original file kept."
+            )
+        except (
+            Exception
+        ) as e:  # Catch other errors like permission issues during rename/remove
+            print(
+                f"ERROR: Error during post-decompression file operations for {gzipped_file_path}: {e}. Original file kept."
+            )
+        finally:
+            if os.path.exists(
+                temp_decompressed_path
+            ):  # Clean up temp file if it still exists
+                os.remove(temp_decompressed_path)
+        return gzipped_file_path  # Return original path if decompression failed
+
     def get_programs(self) -> List[str]:
         xml = self._get_cloud_versions_xml()
         programs = [model.attrib["nome"] for model in xml]
@@ -1057,6 +1362,24 @@ def _budget_matches_list(budget_part: str, all_budgets: List[str]) -> List[str]:
     return [budget for budget in all_budgets if lowered_budget_part in budget.lower()]
 
 
+def _filter_upload_files(directory: str, upload_filter: str) -> List[str]:
+    """
+    Filter files in a directory based on the upload filter.
+    :param directory: Directory to filter files from.
+    :param upload_filter: Regular expression filter for file names.
+    :return: List of filtered file paths.
+    """
+    if not os.path.exists(directory):
+        raise CloudInputError(f"Directory {directory} does not exist")
+
+    regex = re.compile(upload_filter)
+    filtered_files = []
+    for file in os.listdir(directory):
+        if regex.match(file):
+            filtered_files.append(os.path.join(directory, file))
+    return filtered_files
+
+
 def replace_case_str_values(client: Client, case: Case) -> Case:
     """Create a new case object using internal integer IDs instead of string values."""
     # Model Version

psr/factory/__init__.py

@@ -2,6 +2,6 @@
 # Unauthorized copying of this file, via any medium is strictly prohibited
 # Proprietary and confidential
 
-__version__ = "4.1.0b11"
+__version__ = "5.0.0b1"
 
 from .api import *

psr/factory/api.py

@@ -1234,6 +1234,11 @@ class Context(DataObject):
         context_obj._hdr = None
         return context
 
+    def __repr__(self):
+        properties = self.as_dict()
+        props_str = ', '.join(f"{key}={repr(value)}" for key, value in properties.items())
+        return f"psr.factory.Context({props_str})"
+
 
 class Study(_BaseObject):
     def __init__(self):