Stop open's inner Env from clobbering outer env creds

Resolves #1075
This commit is contained in:
Sean Gillies 2017-10-09 12:04:04 -06:00
parent 146e2e1a6a
commit c443affdd6
5 changed files with 129 additions and 64 deletions

View File

@ -253,11 +253,8 @@ def open(fp, mode='r', driver=None, width=None, height=None, count=None,
_, _, scheme = parse_path(fp)
with Env() as env:
# Get AWS credentials only if we're attempting to access a
# raster using the S3 scheme.
if scheme == 's3':
env.get_aws_credentials()
log.debug("AWS credentials have been obtained")
env.credentialize()
# Create dataset instances and pass the given env, which will
# be taken over by the dataset's context manager if it is not

View File

@ -184,6 +184,9 @@ cdef class ConfigEnv(object):
del_gdal_config(key)
log.debug("Unset option %s in env %r", key, self)
def get_config_options(self):
return {k: get_gdal_config(k) for k in self.options}
cdef class GDALEnv(ConfigEnv):
"""Configuration and driver management"""

View File

@ -4,6 +4,7 @@ from functools import wraps
import logging
import threading
import rasterio
from rasterio._env import (
GDALEnv, del_gdal_config, get_gdal_config, set_gdal_config)
from rasterio.compat import string_types
@ -87,7 +88,7 @@ class Env(object):
for GDAL as needed.
"""
def __init__(self, aws_session=None, aws_access_key_id=None,
def __init__(self, session=None, aws_access_key_id=None,
aws_secret_access_key=None, aws_session_token=None,
region_name=None, profile_name=None, **options):
"""Create a new GDAL/AWS environment.
@ -97,7 +98,7 @@ class Env(object):
Parameters
----------
aws_session: object, optional
session: object, optional
A boto3 session.
aws_access_key_id: string, optional
An access key id, as per boto3.
@ -131,38 +132,54 @@ class Env(object):
self.aws_session_token = aws_session_token
self.region_name = region_name
self.profile_name = profile_name
self.aws_session = aws_session
self._creds = (
self.aws_session._session.get_credentials()
if self.aws_session else None)
self.session = session
self.options = options.copy()
self.context_options = {}
def get_aws_credentials(self):
"""Get credentials and configure GDAL."""
import boto3
options = {}
if not self.aws_session:
self.aws_session = boto3.Session(
aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
profile_name=self.profile_name)
self._creds = self.aws_session._session.get_credentials()
self._creds = None
# Pass these credentials to the GDAL environment.
if self._creds.access_key: # pragma: no branch
options.update(aws_access_key_id=self._creds.access_key)
if self._creds.secret_key: # pragma: no branch
options.update(aws_secret_access_key=self._creds.secret_key)
if self._creds.token:
options.update(aws_session_token=self._creds.token)
if self.aws_session.region_name:
options.update(aws_region=self.aws_session.region_name)
@property
def is_credentialized(self):
return bool(self._creds)
# Pass these credentials to the GDAL environment.
local._env.update_config_options(**options)
def credentialize(self):
"""Get credentials and configure GDAL
Note well: this method is a no-op if the GDAL environment
already has credentials, unless session is not None.
"""
if not hascreds():
import boto3
if not self.session and not self.aws_access_key_id and not self.profile_name:
self.session = boto3.Session()
elif not self.session:
self.session = boto3.Session(
aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key,
aws_session_token=self.aws_session_token,
region_name=self.region_name,
profile_name=self.profile_name)
else:
# use self.session
pass
self._creds = self.session._session.get_credentials()
# Pass these credentials to the GDAL environment.
cred_opts = {}
if self._creds.access_key: # pragma: no branch
cred_opts['AWS_ACCESS_KEY_ID'] = self._creds.access_key
if self._creds.secret_key: # pragma: no branch
cred_opts['AWS_SECRET_ACCESS_KEY'] = self._creds.secret_key
if self._creds.token:
cred_opts['AWS_SESSION_TOKEN'] = self._creds.token
if self.session.region_name:
cred_opts['AWS_REGION'] = self.session.region_name
self.options.update(**cred_opts)
setenv(**cred_opts)
def can_credentialize_on_enter(self):
return bool(self.session or self.aws_access_key_id or self.profile_name)
def drivers(self):
"""Return a mapping of registered drivers."""
@ -170,7 +187,6 @@ class Env(object):
def __enter__(self):
log.debug("Entering env context: %r", self)
# No parent Rasterio environment exists.
if local._env is None:
log.debug("Starting outermost env")
@ -194,6 +210,10 @@ class Env(object):
self._has_parent_env = True
self.context_options = getenv()
setenv(**self.options)
if self.can_credentialize_on_enter():
self.credentialize()
log.debug("Entered env context: %r", self)
return self
@ -243,6 +263,10 @@ def getenv():
return local._env.options.copy()
def hasenv():
return bool(local._env)
def setenv(**options):
"""Set options in the existing environment."""
if not local._env:
@ -252,6 +276,12 @@ def setenv(**options):
log.debug("Updated existing %r with options %r", local._env, options)
def hascreds():
gdal_config = local._env.get_config_options()
return bool('AWS_ACCESS_KEY_ID' in gdal_config and
'AWS_SECRET_ACCESS_KEY' in gdal_config)
def delenv():
"""Delete options in the existing environment."""
if not local._env:

View File

@ -21,7 +21,6 @@ def env(ctx, key):
for k, v in sorted(env.drivers().items()):
click.echo("{0}: {1}".format(k, v))
elif key == 'credentials':
env.get_aws_credentials()
click.echo(json.dumps({
'aws_access_key_id': env._creds.access_key,
'aws_secret_access_key': env._creds.secret_key,

View File

@ -10,6 +10,7 @@ import pytest
import rasterio
from rasterio._env import del_gdal_config, get_gdal_config, set_gdal_config
from rasterio._err import CPLE_BaseError
from rasterio.env import defenv, delenv, getenv, setenv, ensure_env
from rasterio.env import default_options
from rasterio.errors import EnvError, RasterioIOError
@ -29,6 +30,7 @@ credentials = pytest.mark.skipif(
logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
L8TIF = "s3://landsat-pds/L8/139/045/LC81390452014295LGN00/LC81390452014295LGN00_B1.TIF"
L8TIFB2 = "s3://landsat-pds/L8/139/045/LC81390452014295LGN00/LC81390452014295LGN00_B2.TIF"
httpstif = "https://landsat-pds.s3.amazonaws.com/L8/139/045/LC81390452014295LGN00/LC81390452014295LGN00_B1.TIF"
@ -137,11 +139,11 @@ def test_aws_session(gdalenv):
aws_session = boto3.Session(
aws_access_key_id='id', aws_secret_access_key='key',
aws_session_token='token', region_name='null-island-1')
s = rasterio.env.Env(aws_session=aws_session)
assert s._creds.access_key == 'id'
assert s._creds.secret_key == 'key'
assert s._creds.token == 'token'
assert s.aws_session.region_name == 'null-island-1'
with rasterio.env.Env(session=aws_session) as s:
assert s._creds.access_key == 'id'
assert s._creds.secret_key == 'key'
assert s._creds.token == 'token'
assert s.session.region_name == 'null-island-1'
def test_aws_session_credentials(gdalenv):
@ -149,26 +151,25 @@ def test_aws_session_credentials(gdalenv):
aws_session = boto3.Session(
aws_access_key_id='id', aws_secret_access_key='key',
aws_session_token='token', region_name='null-island-1')
with rasterio.env.Env(aws_session=aws_session) as s:
s.get_aws_credentials()
assert getenv()['aws_access_key_id'] == 'id'
assert getenv()['aws_region'] == 'null-island-1'
assert getenv()['aws_secret_access_key'] == 'key'
assert getenv()['aws_session_token'] == 'token'
with rasterio.env.Env(session=aws_session) as s:
s.credentialize()
assert getenv()['AWS_ACCESS_KEY_ID'] == 'id'
assert getenv()['AWS_REGION'] == 'null-island-1'
assert getenv()['AWS_SECRET_ACCESS_KEY'] == 'key'
assert getenv()['AWS_SESSION_TOKEN'] == 'token'
def test_with_aws_session_credentials(gdalenv):
"""Create an Env with a boto3 session."""
with rasterio.Env(
aws_access_key_id='id', aws_secret_access_key='key',
aws_session_token='token', region_name='null-island-1') as s:
expected = default_options.copy()
assert getenv() == rasterio.env.local._env.options == expected
s.get_aws_credentials()
env = rasterio.Env(
aws_access_key_id='id', aws_secret_access_key='key',
aws_session_token='token', region_name='null-island-1')
expected = default_options.copy()
with env:
expected.update({
'aws_access_key_id': 'id', 'aws_region': 'null-island-1',
'aws_secret_access_key': 'key', 'aws_session_token': 'token'})
assert getenv() == rasterio.env.local._env.options == expected
'AWS_ACCESS_KEY_ID': 'id', 'AWS_REGION': 'null-island-1',
'AWS_SECRET_ACCESS_KEY': 'key', 'AWS_SESSION_TOKEN': 'token'})
assert getenv() == expected
def test_session_env_lazy(monkeypatch, gdalenv):
@ -177,12 +178,12 @@ def test_session_env_lazy(monkeypatch, gdalenv):
monkeypatch.setenv('AWS_SECRET_ACCESS_KEY', 'key')
monkeypatch.setenv('AWS_SESSION_TOKEN', 'token')
with rasterio.Env() as s:
s.get_aws_credentials()
s.credentialize()
assert getenv() == rasterio.env.local._env.options
expected = {
'aws_access_key_id': 'id',
'aws_secret_access_key': 'key',
'aws_session_token': 'token'}
'AWS_ACCESS_KEY_ID': 'id',
'AWS_SECRET_ACCESS_KEY': 'key',
'AWS_SESSION_TOKEN': 'token'}
for k, v in expected.items():
assert getenv()[k] == v
@ -212,8 +213,18 @@ def test_skip_gtiff(gdalenv):
@mingdalversion
@credentials
@pytest.mark.network
def test_s3_open_with_session(gdalenv):
def test_s3_open_with_env(gdalenv):
"""Read from S3 demonstrating lazy credentials."""
with rasterio.Env():
with rasterio.open(L8TIF) as dataset:
assert dataset.count == 1
@mingdalversion
@credentials
@pytest.mark.network
def test_s3_open_with_implicit_env(gdalenv):
"""Read from S3 using default env."""
with rasterio.open(L8TIF) as dataset:
assert dataset.count == 1
@ -221,10 +232,35 @@ def test_s3_open_with_session(gdalenv):
@mingdalversion
@credentials
@pytest.mark.network
def test_s3_open_with_default_session(gdalenv):
"""Read from S3 using default env."""
with rasterio.open(L8TIF) as dataset:
assert dataset.count == 1
def test_env_open_s3(gdalenv):
"""Read using env as context."""
creds = boto3.Session().get_credentials()
with rasterio.Env(aws_access_key_id=creds.access_key,
aws_secret_access_key=creds.secret_key):
with rasterio.open(L8TIF) as dataset:
assert dataset.count == 1
@mingdalversion
@credentials
@pytest.mark.network
def test_env_open_s3_credentials(gdalenv):
"""Read using env as context."""
aws_session = boto3.Session()
with rasterio.Env(aws_session=aws_session):
with rasterio.open(L8TIF) as dataset:
assert dataset.count == 1
@mingdalversion
@credentials
@pytest.mark.network
def test_ensured_env_no_credentializing(gdalenv):
"""open's extra env doesn't override outer env"""
with rasterio.Env(aws_access_key_id='foo',
aws_secret_access_key='bar'):
with pytest.raises(Exception):
rasterio.open(L8TIFB2)
@mingdalversion