def send(self, request, stream=None, timeout=None, verify=None, cert=None, proxies=None): pathname = url_to_path(request.url) resp = Response() resp.status_code = 200 resp.url = request.url try: stats = lstat(pathname) except (IOError, OSError) as exc: resp.status_code = 404 message = { "error": "file does not exist", "path": pathname, "exception": repr(exc), } fh = SpooledTemporaryFile() fh.write(ensure_binary(json.dumps(message))) fh.seek(0) resp.raw = fh resp.close = resp.raw.close else: modified = formatdate(stats.st_mtime, usegmt=True) content_type = guess_type(pathname)[0] or "text/plain" resp.headers = CaseInsensitiveDict({ "Content-Type": content_type, "Content-Length": stats.st_size, "Last-Modified": modified, }) resp.raw = open(pathname, "rb") resp.close = resp.raw.close return resp
def send(self, request, stream=None, timeout=None, verify=None, cert=None, proxies=None): pathname = url_to_path(request.url) resp = Response() resp.status_code = 200 resp.url = request.url try: stats = lstat(pathname) except (IOError, OSError) as exc: resp.status_code = 404 message = { "error": "file does not exist", "path": pathname, "exception": repr(exc), } fh = SpooledTemporaryFile() fh.write(ensure_binary(json.dumps(message))) fh.seek(0) resp.raw = fh resp.close = resp.raw.close else: modified = formatdate(stats.st_mtime, usegmt=True) content_type = guess_type(pathname)[0] or "text/plain" resp.headers = CaseInsensitiveDict({ "Content-Type": content_type, "Content-Length": stats.st_size, "Last-Modified": modified, }) resp.raw = open(pathname, "rb") resp.close = resp.raw.close return resp
def test_log_artifact_gcp_with_headers( self, databricks_artifact_repo, test_file, artifact_path, expected_location ): expected_headers = {header.name: header.value for header in MOCK_HEADERS} mock_response = Response() mock_response.status_code = 200 mock_response.close = lambda: None with mock.patch( DATABRICKS_ARTIFACT_REPOSITORY + "._get_write_credential_infos" ) as write_credential_infos_mock, mock.patch( "mlflow.utils.rest_utils.cloud_storage_http_request" ) as request_mock: mock_credential_info = ArtifactCredentialInfo( signed_uri=MOCK_GCP_SIGNED_URL, type=ArtifactCredentialType.GCP_SIGNED_URL, headers=MOCK_HEADERS, ) write_credential_infos_mock.return_value = [mock_credential_info] request_mock.return_value = mock_response databricks_artifact_repo.log_artifact(test_file.strpath, artifact_path) write_credential_infos_mock.assert_called_with( run_id=MOCK_RUN_ID, paths=[expected_location] ) request_mock.assert_called_with( "put", MOCK_GCP_SIGNED_URL, data=ANY, headers=expected_headers )
def test_log_artifact_aws_with_headers( self, databricks_artifact_repo, test_file, artifact_path, expected_location ): expected_headers = {header.name: header.value for header in MOCK_HEADERS} mock_response = Response() mock_response.status_code = 200 mock_response.close = lambda: None with mock.patch( DATABRICKS_ARTIFACT_REPOSITORY + "._get_write_credentials" ) as write_credentials_mock, mock.patch( "mlflow.utils.rest_utils.cloud_storage_http_request" ) as request_mock: mock_credentials = ArtifactCredentialInfo( signed_uri=MOCK_AWS_SIGNED_URI, type=ArtifactCredentialType.AWS_PRESIGNED_URL, headers=MOCK_HEADERS, ) write_credentials_response_proto = GetCredentialsForWrite.Response( credentials=mock_credentials ) write_credentials_mock.return_value = write_credentials_response_proto request_mock.return_value = mock_response databricks_artifact_repo.log_artifact(test_file.strpath, artifact_path) write_credentials_mock.assert_called_with(MOCK_RUN_ID, expected_location) request_mock.assert_called_with( "put", MOCK_AWS_SIGNED_URI, ANY, headers=expected_headers )
def send(self, request, stream=None, timeout=None, verify=None, cert=None, proxies=None): parsed_url = urlparse.urlparse(request.url) # We only work for requests with a host of localhost if parsed_url.netloc.lower() != "localhost": raise InvalidURL("Invalid URL %r: Only localhost is allowed" % request.url) real_url = urlparse.urlunparse(parsed_url[:1] + ("",) + parsed_url[2:]) pathname = url_to_path(real_url) resp = Response() resp.status_code = 200 resp.url = real_url stats = os.stat(pathname) modified = email.utils.formatdate(stats.st_mtime, usegmt=True) resp.headers = CaseInsensitiveDict({ "Content-Type": mimetypes.guess_type(pathname)[0] or "text/plain", "Content-Length": stats.st_size, "Last-Modified": modified, }) resp.raw = LocalFSResponse(open(pathname, "rb")) resp.close = resp.raw.close return resp
def send(self, request, stream=None, timeout=None, verify=None, cert=None, proxies=None): parsed_url = urlparse.urlparse(request.url) # We only work for requests with a host of localhost if parsed_url.netloc.lower() != "localhost": raise InvalidURL("Invalid URL %r: Only localhost is allowed" % request.url) real_url = urlparse.urlunparse(parsed_url[:1] + ("",) + parsed_url[2:]) pathname = url_to_path(real_url) resp = Response() resp.status_code = 200 resp.url = real_url stats = os.stat(pathname) modified = email.utils.formatdate(stats.st_mtime, usegmt=True) resp.headers = CaseInsensitiveDict({ "Content-Type": mimetypes.guess_type(pathname)[0] or "text/plain", "Content-Length": stats.st_size, "Last-Modified": modified, }) resp.raw = LocalFSResponse(open(pathname, "rb")) resp.close = resp.raw.close return resp
def test_log_artifact_azure_with_headers(self, databricks_artifact_repo, test_file, artifact_path, expected_location): mock_azure_headers = { "x-ms-encryption-scope": "test-scope", "x-ms-tags": "some-tags", "x-ms-blob-type": "some-type", } filtered_azure_headers = { "x-ms-encryption-scope": "test-scope", "x-ms-tags": "some-tags", } mock_response = Response() mock_response.status_code = 200 mock_response.close = lambda: None with mock.patch( DATABRICKS_ARTIFACT_REPOSITORY + "._get_write_credentials" ) as write_credentials_mock, mock.patch( "mlflow.utils.rest_utils.cloud_storage_http_request" ) as request_mock: mock_credentials = ArtifactCredentialInfo( signed_uri=MOCK_AZURE_SIGNED_URI, type=ArtifactCredentialType.AZURE_SAS_URI, headers=[ ArtifactCredentialInfo.HttpHeader(name=header_name, value=header_value) for header_name, header_value in mock_azure_headers.items() ], ) write_credentials_response_proto = GetCredentialsForWrite.Response( credentials=mock_credentials) write_credentials_mock.return_value = write_credentials_response_proto request_mock.return_value = mock_response databricks_artifact_repo.log_artifact(test_file.strpath, artifact_path) write_credentials_mock.assert_called_with(MOCK_RUN_ID, expected_location) request_mock.assert_called_with( "put", MOCK_AZURE_SIGNED_URI + "?comp=blocklist", ANY, headers=filtered_azure_headers, )
def returnc(self, re: Response): for c in re.iter_content(chunk_size=1024): if c: yield c re.close()
def download_file(http_response: Response, fn: str): with open(fn, 'bw') as f: for chunk in http_response.iter_content(chunk_size=4096): if chunk: f.write(chunk) http_response.close()