first go at having requests use streaming for get/put requests

This commit is contained in:
Ethan Moore 2016-05-27 21:40:31 +00:00 committed by rallytime
parent 7bce4ece1a
commit 4a9b23f03f
2 changed files with 37 additions and 12 deletions

View file

@ -203,7 +203,8 @@ def assumed_creds(prov_dict, role_arn, location=None):
def sig4(method, endpoint, params, prov_dict,
aws_api_version=DEFAULT_AWS_API_VERSION, location=None,
product='ec2', uri='/', requesturl=None, data='', role_arn=None):
product='ec2', uri='/', requesturl=None, data='',
role_arn=None, payload_hash=None):
'''
Sign a query against AWS services using Signature Version 4 Signing
Process. This is documented at:
@ -250,7 +251,8 @@ def sig4(method, endpoint, params, prov_dict,
# Create payload hash (hash of the request body content). For GET
# requests, the payload is an empty string ('').
payload_hash = hashlib.sha256(data).hexdigest()
if not payload_hash:
payload_hash = hashlib.sha256(data).hexdigest()
# Combine elements to create create canonical request
canonical_request = '\n'.join((

View file

@ -92,10 +92,12 @@ def query(key, keyid, method='GET', params=None, headers=None,
keyid = salt.utils.aws.IROLE_CODE
data = ''
payload_hash = None
if method == 'PUT':
if local_file:
with salt.utils.fopen(local_file, 'r') as ifile:
data = ifile.read()
payload_hash = salt.file.get_hash(local_file, form='sha256')
#with salt.utils.fopen(local_file, 'r') as ifile:
# data = ifile.read()
if not requesturl:
requesturl = 'https://{0}/{1}'.format(endpoint, path)
@ -110,6 +112,7 @@ def query(key, keyid, method='GET', params=None, headers=None,
location=location,
product='s3',
requesturl=requesturl,
payload_hash=payload_hash,
)
log.debug('S3 Request: {0}'.format(requesturl))
@ -119,12 +122,31 @@ def query(key, keyid, method='GET', params=None, headers=None,
if not data:
data = None
result = requests.request(method,
requesturl,
headers=headers,
data=data,
verify=verify_ssl)
response = result.content
response = None
if method == 'PUT':
if local_file:
with salt.utils.fopen(local_file, 'r') as data:
result = requests.request(method,
requesturl,
headers=headers,
data=data,
verify=verify_ssl,
stream=True)
response = result.content
elif method == 'GET' and not return_bin:
result = requests.request(method,
requesturl,
headers=headers,
data=data,
verify=verify_ssl,
stream=True)
else:
result = requests.request(method,
requesturl,
headers=headers,
data=data,
verify=verify_ssl)
response = result.content
if result.status_code >= 400:
# On error the S3 API response should contain error message
log.debug(' Response content: {0}'.format(response))
@ -168,8 +190,9 @@ def query(key, keyid, method='GET', params=None, headers=None,
# This can be used to save a binary object to disk
if local_file and method == 'GET':
log.debug('Saving to local file: {0}'.format(local_file))
with salt.utils.fopen(local_file, 'w') as out:
out.write(response)
with salt.utils.fopen(local_file, 'wb') as out:
for chunk in result.iter_content(chunk_size=2048):
out.write(chunk)
return 'Saved to local file: {0}'.format(local_file)
# This can be used to return a binary object wholesale