From a2aae6eedcc82110fad73589580f06683ea83b22 Mon Sep 17 00:00:00 2001
From: MartinMikita <martin.mikita@klokantech.com>
Date: Thu, 3 May 2018 19:26:29 +0200
Subject: [PATCH] Added create argument for gcs.GcsConnection.bucket() and
 s3.S3Connection.bucket().

---
 cloudwrapper/gcs.py |  9 +++++++--
 cloudwrapper/s3.py  | 13 ++++++++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/cloudwrapper/gcs.py b/cloudwrapper/gcs.py
index 2909853..5b91ce2 100644
--- a/cloudwrapper/gcs.py
+++ b/cloudwrapper/gcs.py
@@ -11,7 +11,7 @@
 from .base import BaseBucket
 
 try:
-    from gcloud import storage
+    from gcloud import storage, exceptions
 except ImportError:
     from warnings import warn
     install_modules = [
@@ -35,10 +35,15 @@ def __init__(self):
         self.connection = storage.Client()
 
 
-    def bucket(self, name):
+    def bucket(self, name, create=False):
         for _repeat in range(6):
             try:
                 return Bucket(self.connection.get_bucket(name))
+            except (exceptions.NotFound) as e:
+                if create:
+                    self.connection.create_bucket(name)
+                    continue
+                raise
             except (IOError, BadStatusLine) as e:
                 sleep(_repeat * 2 + 1)
                 if e.errno == errno.EPIPE:
diff --git a/cloudwrapper/s3.py b/cloudwrapper/s3.py
index 4169435..86e60b9 100644
--- a/cloudwrapper/s3.py
+++ b/cloudwrapper/s3.py
@@ -8,6 +8,7 @@
 
 try:
     from boto.s3 import connect_to_region, connection
+    from boto.exceptions import S3ResponseError
 except ImportError:
     from warnings import warn
     install_modules = [
@@ -39,9 +40,15 @@ def __init__(self, region, key=None, secret=None, host=None):
                 aws_secret_access_key=secret)
 
 
-    def bucket(self, name):
-        return Bucket(self.connection.get_bucket(name))
-
+    def bucket(self, name, create=False):
+        for _ in range(6):
+            try:
+                return Bucket(self.connection.get_bucket(name))
+            except S3ResponseError as se:
+                if se.status == 404 and create:
+                    self.connection.create_bucket(name)
+                    continue
+                raise
 
 
 class Bucket(object):