From 5e58818043c73487e86287e4b1452c3c5e6b2ebb Mon Sep 17 00:00:00 2001
From: Arik Fraimovich <arik@arikfr.com>
Date: Wed, 6 Jan 2016 15:14:09 +0200
Subject: [PATCH] Fix CLI to work with organizations

---
 redash/cli/data_sources.py |  5 ++-
 redash/cli/users.py        | 63 ++++++++++++++++++++------------------
 2 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/redash/cli/data_sources.py b/redash/cli/data_sources.py
index 9948fb17..b106fa6a 100644
--- a/redash/cli/data_sources.py
+++ b/redash/cli/data_sources.py
@@ -6,6 +6,7 @@ from redash.query_runner import query_runners, validate_configuration
 
 manager = Manager(help="Data sources management commands.")
 
+
 @manager.command
 def list():
     """List currently configured data sources"""
@@ -27,6 +28,7 @@ def validate_data_source_options(type, options):
         print "Error: invalid configuration."
         exit()
 
+
 @manager.command
 def new(name=None, type=None, options=None):
     """Create new data source"""
@@ -82,7 +84,8 @@ def new(name=None, type=None, options=None):
 
     data_source = models.DataSource.create(name=name,
                                            type=type,
-                                           options=options)
+                                           options=options,
+                                           org=models.Organization.get_by_slug('default'))
     print "Id: {}".format(data_source.id)
 
 
diff --git a/redash/cli/users.py b/redash/cli/users.py
index 5959d364..442cbd74 100644
--- a/redash/cli/users.py
+++ b/redash/cli/users.py
@@ -7,7 +7,7 @@ manager = Manager(help="Users management commands. This commands assume single o
 @manager.option('email', help="email address of the user to grant admin to")
 def grant_admin(email):
     try:
-        user = models.User.get_by_email(email)
+        user = models.User.get_by_email_and_org(email, models.Organization.get_by_slug('default'))
 
         user.groups.append('admin')
         user.save()
@@ -17,33 +17,38 @@ def grant_admin(email):
         print "User [%s] not found." % email
 
 
-# TODO(@arikfr): This needs to be updated to the new org/groups scheme to work.
-# @manager.option('email', help="User's email")
-# @manager.option('name', help="User's full name")
-# @manager.option('--admin', dest='is_admin', action="store_true", default=False, help="set user as admin")
-# @manager.option('--google', dest='google_auth', action="store_true", default=False, help="user uses Google Auth to login")
-# @manager.option('--password', dest='password', default=None, help="Password for users who don't use Google Auth (leave blank for prompt).")
-# @manager.option('--groups', dest='groups', default=models.User.DEFAULT_GROUPS, help="Comma seperated list of groups (leave blank for default).")
-# def create(email, name, groups, is_admin=False, google_auth=False, password=None):
-#     print "Creating user (%s, %s)..." % (email, name)
-#     print "Admin: %r" % is_admin
-#     print "Login with Google Auth: %r\n" % google_auth
-#     if isinstance(groups, basestring):
-#         groups= groups.split(',')
-#         groups.remove('') # in case it was empty string
-#
-#     if is_admin:
-#         groups += ['admin']
-#
-#     user = models.User(email=email, name=name, groups=groups)
-#     if not google_auth:
-#         password = password or prompt_pass("Password")
-#         user.hash_password(password)
-#
-#     try:
-#         user.save()
-#     except Exception, e:
-#         print "Failed creating user: %s" % e.message
+@manager.option('email', help="User's email")
+@manager.option('name', help="User's full name")
+@manager.option('--admin', dest='is_admin', action="store_true", default=False, help="set user as admin")
+@manager.option('--google', dest='google_auth', action="store_true", default=False, help="user uses Google Auth to login")
+@manager.option('--password', dest='password', default=None, help="Password for users who don't use Google Auth (leave blank for prompt).")
+@manager.option('--groups', dest='groups', default=None, help="Comma seperated list of groups (leave blank for default).")
+def create(email, name, groups, is_admin=False, google_auth=False, password=None):
+    print "Creating user (%s, %s)..." % (email, name)
+    print "Admin: %r" % is_admin
+    print "Login with Google Auth: %r\n" % google_auth
+
+    org = models.Organization.get_by_slug('default')
+    if isinstance(groups, basestring):
+        groups= groups.split(',')
+        groups.remove('') # in case it was empty string
+        groups = [int(g) for g in groups]
+
+    if groups is None:
+        groups = [models.Group.get(models.Group.name=="default", models.Group.org==org).id]
+
+    if is_admin:
+        groups += [models.Group.get(models.Group.name=="admin", models.Group.org==org).id]
+
+    user = models.User(email=email, name=name, groups=groups)
+    if not google_auth:
+        password = password or prompt_pass("Password")
+        user.hash_password(password)
+
+    try:
+        user.save()
+    except Exception, e:
+        print "Failed creating user: %s" % e.message
 
 
 @manager.option('email', help="email address of user to delete")
@@ -56,7 +61,7 @@ def delete(email):
 @manager.option('email', help="email address of the user to change password for")
 def password(email, password):
     try:
-        user = models.User.get_by_email(email)
+        user = models.User.get_by_email_and_org(email, models.Organization.get_by_slug('default'))
 
         user.hash_password(password)
         user.save()