From a5da8c47bdfb2c3899e7ab4d5534b387032f2a31 Mon Sep 17 00:00:00 2001 From: Denis Rykov Date: Wed, 2 Dec 2020 19:55:14 +0100 Subject: [PATCH] Add support for Microsoft Azure Blob Storage (#1906) Co-authored-by: Sean Gillies --- rasterio/path.py | 3 +- rasterio/session.py | 66 +++++++++++++++++++++++++++++++++++++++++++ tests/test_env.py | 24 +++++++++++++++- tests/test_session.py | 43 +++++++++++++++++++++++++++- 4 files changed, 133 insertions(+), 3 deletions(-) diff --git a/rasterio/path.py b/rasterio/path.py index 9c906732..15247a97 100644 --- a/rasterio/path.py +++ b/rasterio/path.py @@ -21,12 +21,13 @@ SCHEMES = { 'file': 'file', 'oss': 'oss', 'gs': 'gs', + 'az': 'az', } CURLSCHEMES = set([k for k, v in SCHEMES.items() if v == 'curl']) # TODO: extend for other cloud plaforms. -REMOTESCHEMES = set([k for k, v in SCHEMES.items() if v in ('curl', 's3', 'oss', 'gs',)]) +REMOTESCHEMES = set([k for k, v in SCHEMES.items() if v in ('curl', 's3', 'oss', 'gs', 'az',)]) class Path(object): diff --git a/rasterio/session.py b/rasterio/session.py index 977f8828..61532040 100644 --- a/rasterio/session.py +++ b/rasterio/session.py @@ -113,6 +113,9 @@ class Session(object): elif path.path.startswith("/vsiswift/"): return SwiftSession + elif path.scheme == "az": + return AzureSession + # This factory can be extended to other cloud providers here. # elif path.scheme == "cumulonimbus": # for example. # return CumulonimbusSession(*args, **kwargs) @@ -535,3 +538,66 @@ class SwiftSession(Session): dict """ return {k.upper(): v for k, v in self.credentials.items()} + + +class AzureSession(Session): + """Configures access to secured resources stored in Microsoft Azure Blob Storage. + """ + def __init__(self, azure_storage_connection_string=None, + azure_storage_account=None, azure_storage_access_key=None): + """Create new Microsoft Azure Blob Storage session + + Parameters + ---------- + azure_storage_connection_string: string + A connection string contains both an account name and a secret key. + azure_storage_account: string + An account name + azure_storage_access_key: string + A secret key + """ + + if azure_storage_connection_string: + self._creds = { + "azure_storage_connection_string": azure_storage_connection_string + } + else: + self._creds = { + "azure_storage_account": azure_storage_account, + "azure_storage_access_key": azure_storage_access_key + } + + @classmethod + def hascreds(cls, config): + """Determine if the given configuration has proper credentials + + Parameters + ---------- + cls : class + A Session class. + config : dict + GDAL configuration as a dict. + + Returns + ------- + bool + + """ + return 'AZURE_STORAGE_CONNECTION_STRING' in config or ( + 'AZURE_STORAGE_ACCOUNT' in config and 'AZURE_STORAGE_ACCESS_KEY' in config + ) + + @property + def credentials(self): + """The session credentials as a dict""" + return self._creds + + def get_credential_options(self): + """Get credentials as GDAL configuration options + + Returns + ------- + dict + + """ + return {k.upper(): v for k, v in self.credentials.items()} diff --git a/tests/test_env.py b/tests/test_env.py index 9780d012..d273b4c9 100644 --- a/tests/test_env.py +++ b/tests/test_env.py @@ -19,7 +19,7 @@ from rasterio.env import Env, defenv, delenv, getenv, setenv, ensure_env, ensure from rasterio.env import GDALVersion, require_gdal_version from rasterio.errors import EnvError, RasterioIOError, GDALVersionError from rasterio.rio.main import main_group -from rasterio.session import AWSSession, DummySession, OSSSession, SwiftSession +from rasterio.session import AWSSession, DummySession, OSSSession, SwiftSession, AzureSession from .conftest import requires_gdal21 @@ -845,6 +845,28 @@ def test_swift_session_credentials(gdalenv): assert getenv()['SWIFT_AUTH_TOKEN'] == 'bar' +def test_azure_session_credentials(gdalenv): + """Create an Env with azure session.""" + azure_session = AzureSession( + azure_storage_account='foo', + azure_storage_access_key='bar' + ) + with rasterio.env.Env(session=azure_session) as s: + s.credentialize() + assert getenv()['AZURE_STORAGE_ACCOUNT'] == 'foo' + assert getenv()['AZURE_STORAGE_ACCESS_KEY'] == 'bar' + + +def test_azure_session_credentials_connection_string(gdalenv): + """Create an Env with azure session.""" + azure_session = AzureSession( + azure_storage_connection_string='AccountName=myaccount;AccountKey=MY_ACCOUNT_KEY', + ) + with rasterio.env.Env(session=azure_session) as s: + s.credentialize() + assert getenv()['AZURE_STORAGE_CONNECTION_STRING'] == 'AccountName=myaccount;AccountKey=MY_ACCOUNT_KEY' + + def test_swift_session_by_user_key(): def mock_init( self, session=None, diff --git a/tests/test_session.py b/tests/test_session.py index f0f3e1de..824798c7 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -7,7 +7,15 @@ try: except ImportError: import mock -from rasterio.session import DummySession, AWSSession, Session, OSSSession, GSSession, SwiftSession +from rasterio.session import ( + DummySession, + AWSSession, + Session, + OSSSession, + GSSession, + SwiftSession, + AzureSession, +) def test_base_session_hascreds_notimpl(): @@ -236,3 +244,36 @@ def test_no_credentialization_if_unsigned(monkeypatch): """Don't get credentials if we're not signing, see #1984""" sesh = AWSSession(aws_unsigned=True) assert sesh._creds is None + + +def test_azure_session_class(): + """AzureSession works""" + azure_session = AzureSession(azure_storage_account='foo', azure_storage_access_key='bar') + assert azure_session._creds + assert azure_session.get_credential_options()['AZURE_STORAGE_ACCOUNT'] == 'foo' + assert azure_session.get_credential_options()['AZURE_STORAGE_ACCESS_KEY'] == 'bar' + + +def test_azure_session_class_connection_string(): + """AzureSession works""" + azure_session = AzureSession(azure_storage_connection_string='AccountName=myaccount;AccountKey=MY_ACCOUNT_KEY') + assert azure_session._creds + assert ( + azure_session.get_credential_options()['AZURE_STORAGE_CONNECTION_STRING'] + == 'AccountName=myaccount;AccountKey=MY_ACCOUNT_KEY' + ) + + +def test_session_factory_az_kwargs(): + """Get an AzureSession for az:// paths with keywords""" + sesh = Session.from_path("az://lol/wut", azure_storage_account='foo', azure_storage_access_key='bar') + assert isinstance(sesh, AzureSession) + assert sesh.get_credential_options()['AZURE_STORAGE_ACCOUNT'] == 'foo' + assert sesh.get_credential_options()['AZURE_STORAGE_ACCESS_KEY'] == 'bar' + + +def test_session_factory_az_kwargs_connection_string(): + """Get an AzureSession for az:// paths with keywords""" + sesh = Session.from_path("az://lol/wut", azure_storage_connection_string='AccountName=myaccount;AccountKey=MY_ACCOUNT_KEY') + assert isinstance(sesh, AzureSession) + assert sesh.get_credential_options()['AZURE_STORAGE_CONNECTION_STRING'] == 'AccountName=myaccount;AccountKey=MY_ACCOUNT_KEY' \ No newline at end of file