Azure integration

lazydata/cli/commands/addremote.py

@@ -6,21 +6,29 @@
 from lazydata.storage.remote import RemoteStorage
 
 import lazy_import
+
 botocore = lazy_import.lazy_module("botocore")
 
-class AddRemoteCommand(BaseCommand):
 
+class AddRemoteCommand(BaseCommand):
     def add_arguments(self, parser):
-        parser.add_argument('url', type=str, help='URL of the remote storage backend')
-        parser.add_argument('--endpoint-url', type=str, nargs='?', help='The complete URL to use for the constructed client')
+        parser.add_argument("url", type=str, help="URL of the remote storage backend")
+        parser.add_argument(
+            "--endpoint-url",
+            type=str,
+            nargs="?",
+            help="The complete URL to use for the constructed client",
+        )
         return parser
 
     def handle(self, args):
         url = args.url
         endpoint_url = args.endpoint_url
 
-        if not url.startswith("s3://"):
-            print("ERROR: Only S3 URLs are currently supported. For example: `s3://mybucket` or `s3://mybucket/myfolder`")
+        if not url.startswith(("s3://", "az://")):
+            print(
+                "ERROR: Only S3 URLs are currently supported. For example: `s3://mybucket` or `s3://mybucket/myfolder`"
+            )
             sys.exit(1)
 
         success = False
@@ -33,17 +41,25 @@ def handle(self, args):
                     success = True
                 else:
                     success = True
-                    print("ERROR: The remote storage location you specified does not exist or is not accessible to you")
+                    print(
+                        "ERROR: The remote storage location you specified does not exist or is not accessible to you"
+                    )
             except botocore.exceptions.NoCredentialsError:
                 success = False
+            # Azure does not throw a credentials error. There does not seem
+            # to be a way to differentiate exists vs credentials with Azure.
 
             if not success:
                 print("ERROR: No valid AWS credentials found.")
                 config_now = input("Do you want to configure AWS credentials now? [y/n] ")
                 if config_now.strip() == "y":
                     setup_aws_credentials()
-                    print("Credentials successfully stored. Trying again with these new credentials...")
+                    print(
+                        "Credentials successfully stored. Trying again with these new credentials..."
+                    )
                 else:
                     success = True
-                    print("Alright, will not configure credentials now. Re-run this command to try again, "
-                          "or configure using the AWS CLI: `aws configure`.")
+                    print(
+                        "Alright, will not configure credentials now. Re-run this command to try again, "
+                        "or configure using the AWS CLI: `aws configure`."
+                    )

lazydata/cli/commands/config.py

@@ -1,14 +1,11 @@
 from lazydata.cli.commands.BaseCommand import BaseCommand
-from lazydata.storage.cloudsetup import setup_aws_credentials
+from lazydata.storage.cloudsetup import setup_credentials
 
 
 class ConfigCommand(BaseCommand):
     def add_arguments(self, parser):
-        parser.add_argument('backend', type=str, help='The backend to configure, currently only `aws` supported')
+        parser.add_argument('backend', type=str, help='The backend to configure, currently only `aws` and `azure` are supported')
         return parser
 
     def handle(self, args):
-        if args.backend == "aws":
-            setup_aws_credentials()
-        else:
-            print("ERROR: Unrecognised backend `%s`. Currently supported: `aws`" % args.backend)
+        setup_credentials(args.backend)

lazydata/storage/cloudsetup.py

@@ -2,8 +2,22 @@
 Common functions to interactively set up credentials for cloud access
 """
 
+import configparser
 from pathlib import Path
 
+azure_root = Path(Path.home(), ".azure")
+azure_config_path = Path(azure_root, "config")
+
+
+def setup_credentials(backend: str):
+    if backend == "aws":
+        setup_aws_credentials()
+    elif backend == "azure":
+        setup_azure_credentials()
+    else:
+        print("ERROR: Unrecognised backend `%s`. Currently supported: `aws` and `azure`" % backend)
+
+
 def setup_aws_credentials():
 
     public_key = input("Your AWS access key: ")
@@ -20,17 +34,47 @@ def setup_aws_credentials():
     if not aws_root.exists():
         aws_root.mkdir()
 
+    aws_cred_parser = configparser.ConfigParser()
+    aws_config_parser = configparser.ConfigParser()
+
+    if aws_cred.exists():
+        with open(str(aws_cred), "r") as f:
+            aws_cred_parser.read_file(f)
+
+    aws_cred_parser.set("default", "aws_access_key_id", public_key)
+    aws_cred_parser.set("default", "aws_secret_access_key", private_key)
+
     with open(str(aws_cred), "w") as f:
-        f.writelines([
-            "[default]\n",
-            "aws_access_key_id = %s\n" % public_key,
-            "aws_secret_access_key = %s\n" % private_key,
-        ])
+        aws_cred_parser.write(f)
+
+    if aws_config.exists():
+        with open(str(aws_config), "r") as f:
+            aws_config_parser.read_file(f)
+
+    aws_config_parser.set("default", "region", default_zone)
 
     with open(str(aws_config), "w") as f:
-        f.writelines([
-            "[default]\n",
-            "region = %s\n" % default_zone,
-        ])
+        aws_config_parser.write(f)
 
     print("Credentials written into %s" % str(aws_root))
+
+
+def setup_azure_credentials():
+
+    private_key = input("Your Azure storage access key: ")
+
+    if not azure_root.exists():
+        azure_root.mkdir()
+
+    azure_config_parser = configparser.ConfigParser()
+
+    if azure_config_path.exists():
+        with open(str(azure_config_path), "r") as f:
+            azure_config_parser.read_file(f)
+
+    azure_config_parser.set("storage", "key", private_key)
+
+    with open(str(azure_config_path), "w") as f:
+        azure_config_parser.write(f)
+
+    print("Credentials written into %s" % str(azure_root))