Sfoglia il codice sorgente

HUE-8530 [organization] Port User and Group model switch everywhere

Romain 6 anni fa
parent
commit
672df534b9
82 ha cambiato i file con 594 aggiunte e 217 eliminazioni
  1. 5 1
      apps/about/src/about/tests.py
  2. 9 6
      apps/beeswax/src/beeswax/api.py
  3. 6 2
      apps/beeswax/src/beeswax/management/commands/beeswax_install_examples.py
  4. 6 2
      apps/beeswax/src/beeswax/test_base.py
  5. 8 2
      apps/beeswax/src/beeswax/tests.py
  6. 6 2
      apps/beeswax/src/beeswax/views.py
  7. 9 2
      apps/filebrowser/src/filebrowser/forms.py
  8. 8 3
      apps/filebrowser/src/filebrowser/views.py
  9. 20 15
      apps/filebrowser/src/filebrowser/views_test.py
  10. 6 1
      apps/hbase/src/hbase/management/commands/hbase_setup.py
  11. 6 2
      apps/hbase/src/hbase/tests.py
  12. 6 3
      apps/impala/src/impala/tests.py
  13. 6 1
      apps/jobbrowser/src/jobbrowser/tests.py
  14. 0 1
      apps/jobsub/src/jobsub/models.py
  15. 7 1
      apps/jobsub/src/jobsub/tests.py
  16. 6 1
      apps/metastore/src/metastore/tests.py
  17. 6 2
      apps/oozie/src/oozie/models2_tests.py
  18. 7 1
      apps/oozie/src/oozie/tests.py
  19. 6 1
      apps/pig/src/pig/tests.py
  20. 6 1
      apps/search/src/search/tests.py
  21. 6 1
      apps/security/src/security/tests.py
  22. 7 2
      apps/sqoop/src/sqoop/tests.py
  23. 6 2
      apps/useradmin/src/useradmin/api.py
  24. 6 7
      apps/useradmin/src/useradmin/forms.py
  25. 10 6
      apps/useradmin/src/useradmin/ldap_access.py
  26. 6 2
      apps/useradmin/src/useradmin/middleware.py
  27. 6 1
      apps/useradmin/src/useradmin/test_ldap.py
  28. 7 1
      apps/useradmin/src/useradmin/test_ldap_deprecated.py
  29. 6 2
      apps/useradmin/src/useradmin/tests.py
  30. 6 1
      apps/useradmin/src/useradmin/tests_api.py
  31. 10 8
      apps/useradmin/src/useradmin/views.py
  32. 6 2
      desktop/core/src/desktop/api2_tests.py
  33. 6 3
      desktop/core/src/desktop/api_tests.py
  34. 5 1
      desktop/core/src/desktop/auth/views.py
  35. 6 1
      desktop/core/src/desktop/auth/views_test.py
  36. 7 2
      desktop/core/src/desktop/configuration/tests.py
  37. 12 7
      desktop/core/src/desktop/converter_tests.py
  38. 7 2
      desktop/core/src/desktop/lib/django_test_util.py
  39. 9 3
      desktop/core/src/desktop/lib/fs/proxyfs.py
  40. 6 2
      desktop/core/src/desktop/lib/fs/proxyfs_test.py
  41. 7 1
      desktop/core/src/desktop/management/commands/convert_documents.py
  42. 6 1
      desktop/core/src/desktop/management/commands/create_user_directories.py
  43. 6 2
      desktop/core/src/desktop/metrics.py
  44. 6 2
      desktop/core/src/desktop/middleware.py
  45. 6 1
      desktop/core/src/desktop/monkey_patches.py
  46. 9 2
      desktop/core/src/desktop/tests.py
  47. 6 1
      desktop/core/src/desktop/tests_doc2.py
  48. 7 2
      desktop/libs/aws/src/aws/s3/s3fs_test.py
  49. 6 1
      desktop/libs/dashboard/src/dashboard/tests.py
  50. 6 1
      desktop/libs/indexer/src/indexer/indexers/envelope.py
  51. 42 36
      desktop/libs/indexer/src/indexer/indexers/envelope_tests.py
  52. 6 1
      desktop/libs/indexer/src/indexer/indexers/flume.py
  53. 8 3
      desktop/libs/indexer/src/indexer/indexers/flume_tests.py
  54. 7 3
      desktop/libs/indexer/src/indexer/indexers/morphline.py
  55. 7 2
      desktop/libs/indexer/src/indexer/indexers/morphline_tests.py
  56. 7 2
      desktop/libs/indexer/src/indexer/indexers/rdbms.py
  57. 7 1
      desktop/libs/indexer/src/indexer/indexers/sql.py
  58. 6 1
      desktop/libs/indexer/src/indexer/indexers/sql_tests.py
  59. 8 3
      desktop/libs/indexer/src/indexer/rdbms_indexer_tests.py
  60. 6 1
      desktop/libs/indexer/src/indexer/solr_client_tests.py
  61. 6 1
      desktop/libs/indexer/src/indexer/tests.py
  62. 15 11
      desktop/libs/liboauth/src/liboauth/backend.py
  63. 9 4
      desktop/libs/liboauth/src/liboauth/views.py
  64. 6 1
      desktop/libs/liboozie/src/liboozie/submittion2_tests.py
  65. 6 1
      desktop/libs/liboozie/src/liboozie/submittion_tests.py
  66. 6 1
      desktop/libs/liboozie/src/liboozie/tests.py
  67. 6 1
      desktop/libs/libopenid/src/libopenid/backend.py
  68. 7 2
      desktop/libs/libsaml/src/libsaml/backend.py
  69. 6 1
      desktop/libs/libsentry/src/libsentry/test_privilege_checker.py
  70. 6 2
      desktop/libs/libsentry/src/libsentry/tests.py
  71. 6 1
      desktop/libs/libsolr/src/libsolr/tests.py
  72. 6 2
      desktop/libs/libzookeeper/src/libzookeeper/tests.py
  73. 6 1
      desktop/libs/metadata/src/metadata/catalog/navigator_client_tests.py
  74. 5 1
      desktop/libs/metadata/src/metadata/catalog_tests.py
  75. 6 2
      desktop/libs/metadata/src/metadata/optimizer_api_tests.py
  76. 6 2
      desktop/libs/metadata/src/metadata/optimizer_client_tests.py
  77. 6 2
      desktop/libs/notebook/src/notebook/connectors/hiveserver2_tests.py
  78. 6 2
      desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py
  79. 6 1
      desktop/libs/notebook/src/notebook/management/commands/notebook_setup.py
  80. 7 2
      desktop/libs/notebook/src/notebook/models.py
  81. 7 1
      desktop/libs/notebook/src/notebook/tasks.py
  82. 6 2
      desktop/libs/notebook/src/notebook/tests.py

+ 5 - 1
apps/about/src/about/tests.py

@@ -18,16 +18,20 @@
 from builtins import object
 import json
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from nose.tools import assert_true, assert_false, assert_equal
 
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Settings
 from oozie.tests import OozieBase
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
 
 
 class TestAboutBase(object):

+ 9 - 6
apps/beeswax/src/beeswax/api.py

@@ -20,13 +20,14 @@ import logging
 import json
 import re
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.http import Http404
 from django.utils.translation import ugettext as _
 from django.views.decorators.http import require_POST
 
 from thrift.transport.TTransport import TTransportException
+from desktop.auth.backend import is_admin
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.context_processors import get_app_name
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.exceptions import StructuredThriftTransportException
@@ -37,7 +38,6 @@ from metastore import parser
 from notebook.models import escape_rows, MockedDjangoRequest, make_notebook
 
 import beeswax.models
-
 from beeswax.data_export import upload
 from beeswax.design import HQLdesign
 from beeswax.conf import USE_GET_LOG_API
@@ -45,14 +45,17 @@ from beeswax.forms import QueryForm
 from beeswax.models import Session, QueryHistory
 from beeswax.server import dbms
 from beeswax.server.dbms import expand_exception, get_query_server_config, QueryServerException, QueryServerTimeoutException,\
-  SubQueryTable
+    SubQueryTable
 from beeswax.views import authorized_get_design, authorized_get_query_history, make_parameterization_form,\
-                          safe_get_design, save_design, massage_columns_for_json, _get_query_handle_and_state, \
-                          parse_out_jobs
+    safe_get_design, save_design, massage_columns_for_json, _get_query_handle_and_state, parse_out_jobs
 from metastore.conf import FORCE_HS2_METADATA
 from metastore.views import _get_db, _get_servername
 
-from desktop.auth.backend import is_admin
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
apps/beeswax/src/beeswax/management/commands/beeswax_install_examples.py

@@ -23,11 +23,10 @@ import os
 import pwd
 
 from django.core.management.base import BaseCommand
-from django.contrib.auth.models import User
 from django.utils.translation import ugettext as _
 
 from desktop.lib.exceptions_renderable import PopupException
-from desktop.conf import USE_NEW_EDITOR
+from desktop.conf import USE_NEW_EDITOR, ENABLE_ORGANIZATIONS
 from desktop.models import Directory, Document, Document2, Document2Permission
 from hadoop import cluster
 from notebook.models import import_saved_beeswax_query
@@ -40,6 +39,11 @@ from beeswax.hive_site import has_concurrency_support
 from beeswax.server import dbms
 from beeswax.server.dbms import get_query_server_config, QueryServerException
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
apps/beeswax/src/beeswax/test_base.py

@@ -31,8 +31,8 @@ import time
 
 from nose.tools import assert_true, assert_false
 from django.urls import reverse
-from django.contrib.auth.models import User
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.paths import get_run_root
 from desktop.lib.python_util import find_unused_port
@@ -43,10 +43,14 @@ from hadoop import cluster, pseudo_hdfs4
 from hadoop.pseudo_hdfs4 import is_live_cluster, get_db_prefix
 
 import beeswax.conf
-
 from beeswax.server.dbms import get_query_server_config
 from beeswax.server import dbms
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 HIVE_SERVER_TEST_PORT = find_unused_port()
 _INITIALIZED = False

+ 8 - 2
apps/beeswax/src/beeswax/tests.py

@@ -44,7 +44,6 @@ from nose.plugins.skip import SkipTest
 
 from django.utils.encoding import smart_str
 from django.utils.html import escape
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.db import transaction
 
@@ -55,7 +54,7 @@ from desktop.conf import \
     AUTH_PASSWORD_SCRIPT as DEFAULT_AUTH_PASSWORD_SCRIPT, \
     LDAP_USERNAME, \
     LDAP_PASSWORD, \
-    USE_NEW_EDITOR
+    USE_NEW_EDITOR, ENABLE_ORGANIZATIONS
 from desktop import redaction
 from desktop.redaction import logfilter
 from desktop.redaction.engine import RedactionPolicy, RedactionRule
@@ -98,8 +97,15 @@ if sys.version_info[0] > 2:
 else:
   from cStringIO import StringIO as string_io
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
+
 def _list_dir_without_temp_files(fs, target_dir):
   return [f for f in fs.listdir(target_dir) if not f.startswith('.')]
 

+ 6 - 2
apps/beeswax/src/beeswax/views.py

@@ -26,7 +26,6 @@ import time
 from django import forms
 from django.core.paginator import Paginator, EmptyPage, InvalidPage
 from django.contrib import messages
-from django.contrib.auth.models import User
 from django.db.models import Q
 from django.http import HttpResponse, QueryDict
 from django.shortcuts import redirect
@@ -35,7 +34,7 @@ from django.utils.translation import ugettext as _
 from django.urls import reverse
 
 from desktop.appmanager import get_apps_dict
-from desktop.conf import ENABLE_DOWNLOAD, REDIRECT_WHITELIST
+from desktop.conf import ENABLE_DOWNLOAD, REDIRECT_WHITELIST, ENABLE_ORGANIZATIONS
 from desktop.context_processors import get_app_name
 
 from desktop.lib.django_util import JsonResponse
@@ -58,6 +57,11 @@ from beeswax.server.dbms import expand_exception, get_query_server_config, Query
 
 from desktop.auth.backend import is_admin
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 9 - 2
apps/filebrowser/src/filebrowser/forms.py

@@ -24,23 +24,30 @@ import sys
 import urllib.request, urllib.error
 
 from django import forms
-from django.contrib.auth.models import User, Group
 from django.forms import FileField, CharField, BooleanField, Textarea
 from django.forms.formsets import formset_factory, BaseFormSet
+from django.utils.translation import ugettext_lazy as _
 
 from aws.s3 import S3A_ROOT, normpath as s3_normpath
 from azure.abfs.__init__ import ABFS_ROOT, normpath as abfs_normpath
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib import i18n
 from hadoop.fs import normpath
+
 from filebrowser.lib import rwx
 
-from django.utils.translation import ugettext_lazy as _
+
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group, default_organization
+else:
+  from django.contrib.auth.models import User, Group
 
 if sys.version_info[0] > 2:
   from urllib.parse import unquote as urllib_unquote
 else:
   from urllib import unquote as urllib_unquote
 
+
 logger = logging.getLogger(__name__)
 
 

+ 8 - 3
apps/filebrowser/src/filebrowser/views.py

@@ -35,7 +35,6 @@ from bz2 import decompress
 from datetime import datetime
 from gzip import GzipFile
 
-from django.contrib.auth.models import User, Group
 from django.core.paginator import EmptyPage, Paginator, Page, InvalidPage
 from django.urls import reverse
 from django.template.defaultfilters import stringformat, filesizeformat
@@ -51,6 +50,7 @@ from django.utils.translation import ugettext as _
 from aws.s3.s3fs import S3FileSystemException, S3ListAllBucketsException
 from desktop import appmanager
 from desktop.auth.backend import is_admin
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib import i18n
 from desktop.lib.conf import coerce_bool
 from desktop.lib.django_util import render, format_preserving_redirect
@@ -74,10 +74,15 @@ from filebrowser.lib.archives import archive_factory
 from filebrowser.lib.rwx import filetype, rwx
 from filebrowser.lib import xxd
 from filebrowser.forms import RenameForm, UploadFileForm, UploadArchiveForm, MkDirForm, EditorForm, TouchForm,\
-                              RenameFormSet, RmTreeFormSet, ChmodFormSet, ChownFormSet, CopyFormSet, RestoreFormSet,\
-                              TrashPurgeForm, SetReplicationFactorForm
+    RenameFormSet, RmTreeFormSet, ChmodFormSet, ChownFormSet, CopyFormSet, RestoreFormSet,\
+    TrashPurgeForm, SetReplicationFactorForm
 
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 if sys.version_info[0] > 2:
   import io
   from io import StringIO as string_io

+ 20 - 15
apps/filebrowser/src/filebrowser/views_test.py

@@ -39,15 +39,17 @@ from avro import schema, datafile, io
 
 from aws.s3.s3fs import S3FileSystemException
 from aws.s3.s3test_utils import get_test_bucket
+
 from azure.conf import is_abfs_enabled, is_adls_enabled
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.utils.encoding import smart_str
+from django.utils.translation import ugettext_lazy as _
 from nose.plugins.attrib import attr
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises,\
   assert_greater
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_to_group, add_permission, remove_from_group
 from desktop.lib.view_util import location_to_url
@@ -59,13 +61,16 @@ from filebrowser.lib.rwx import expand_mode
 from filebrowser.views import snappy_installed
 from desktop.conf import is_oozie_enabled
 
-
-
 if sys.version_info[0] > 2:
   from urllib.parse import unquote as urllib_unquote
 else:
   from urllib import unquote as urllib_unquote
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group, default_organization
+else:
+  from django.contrib.auth.models import User, Group
+
 
 LOG = logging.getLogger(__name__)
 
@@ -982,7 +987,7 @@ alert("XSS")
   def test_compress_hdfs_files(self):
     if not is_oozie_enabled():
       raise SkipTest
-    
+
     def make_and_test_dir(pre, test_direct):
       test_dir = pre + "/" + test_direct
       test_file = test_dir + '/test.txt'
@@ -992,12 +997,12 @@ alert("XSS")
       for i in range(3):
         f = self.cluster.fs.open(test_file + "%s" %i, "w")
         f.close()
-        
+
       resp = self.c.post('/filebrowser/compress_files', {'upload_path': pre, 'files[]': [test_direct], 'archive_name': 'test_compress.zip'})
       response = json.loads(resp.content)
       assert_equal(0, response['status'], response)
       assert_true('handle' in response and response['handle']['id'], response)
-      responseid = '"' + response['handle']['id'] + '"' 
+      responseid = '"' + response['handle']['id'] + '"'
       timeout_time = time() + 25
       end_time = time()
       while timeout_time > end_time:
@@ -1005,16 +1010,16 @@ alert("XSS")
         response2 = json.loads(resp2.content)
         if response2['app']['status'] != 'RUNNING':
           assert_equal(response2['app']['status'] , 'SUCCEEDED', response2)
-          break 
+          break
         sleep(3)
         end_time = time()
       assert_greater(timeout_time, end_time, response)
-        
-        
+
+
     ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(True)
     prefix = self.cluster.fs_prefix + '/test_compress_files'
     self.cluster.fs.mkdir(prefix)
-    
+
     try:
       make_and_test_dir(prefix, 'testdir')
       make_and_test_dir(prefix, 'test dir1')
@@ -1023,7 +1028,7 @@ alert("XSS")
     finally:
       ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(False)
       cleanup_tree(self.cluster, prefix)
-      
+
 
   def test_extract_tgz(self):
     ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(True)
@@ -1243,7 +1248,7 @@ class TestS3AccessPermissions(object):
       assert_equal(200, response.status_code)
     finally:
       remove_from_group(self.user.username, 'has_s3')
-      
+
 class TestABFSAccessPermissions(object):
 
   def setUp(self):
@@ -1275,7 +1280,7 @@ class TestABFSAccessPermissions(object):
       assert_equal(200, response.status_code)
     finally:
       remove_from_group(self.user.username, 'has_abfs')
-      
+
 class TestADLSAccessPermissions(object):
 
   def setUp(self):
@@ -1290,7 +1295,7 @@ class TestADLSAccessPermissions(object):
       raise SkipTest
     response = self.client.get('/filebrowser/view=ADL://')
     assert_equal(500, response.status_code)
-    
+
     response = self.client.get('/filebrowser/view=ADL://hue_adls_testing')
     assert_equal(500, response.status_code)
 
@@ -1321,4 +1326,4 @@ class TestADLSAccessPermissions(object):
       response = self.client.get('/filebrowser/view=ADL://')
       assert_equal(200, response.status_code)
     finally:
-      remove_from_group(self.user.username, 'has_adls')
+      remove_from_group(self.user.username, 'has_adls')

+ 6 - 1
apps/hbase/src/hbase/management/commands/hbase_setup.py

@@ -20,10 +20,10 @@ import os
 
 from datetime import datetime, timedelta
 
-from django.contrib.auth.models import User
 from django.core.management.base import BaseCommand
 from django.utils.translation import ugettext as _
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.paths import get_apps_root
 
 from useradmin.models import install_sample_user
@@ -31,6 +31,11 @@ from useradmin.models import install_sample_user
 from hbased.ttypes import AlreadyExists
 from hbase.api import HbaseApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
apps/hbase/src/hbase/tests.py

@@ -23,8 +23,7 @@ import tempfile
 
 from nose.tools import assert_true, assert_equal
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_to_group
 
@@ -35,6 +34,11 @@ from hadoop.pseudo_hdfs4 import is_live_cluster
 from nose.plugins.skip import SkipTest
 from django.urls import reverse
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 def test_security_plain():
   tmpdir = tempfile.mkdtemp()

+ 6 - 3
apps/impala/src/impala/tests.py

@@ -24,10 +24,10 @@ from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false, assert_raises
 from mock import patch, Mock
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 import desktop.conf as desktop_conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.test_utils import add_to_group
@@ -36,7 +36,6 @@ from hadoop.pseudo_hdfs4 import get_db_prefix, is_live_cluster
 
 from beeswax import data_export
 from beeswax.design import hql_query
-
 from beeswax.data_export import download
 from beeswax.models import SavedQuery, QueryHistory
 from beeswax.server import dbms
@@ -46,6 +45,11 @@ from beeswax.tests import _make_query
 from impala import conf
 from impala.dbms import ImpalaDbms
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 
@@ -516,4 +520,3 @@ class TestImpalaDbms(object):
                  ('order_date', '`default`.`customers`.`orders`'))
     assert_equal(ImpalaDbms.get_nested_select('default', 'customers', 'orders', 'item/items/item/product_id'),
                  ('product_id', '`default`.`customers`.`orders`.`items`'))
-

+ 6 - 1
apps/jobbrowser/src/jobbrowser/tests.py

@@ -24,11 +24,11 @@ import re
 import time
 import unittest
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_false, assert_equal, assert_raises
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_to_group
 from desktop.models import Document
@@ -48,6 +48,11 @@ from jobbrowser.conf import SHARE_JOBS
 from jobbrowser.models import can_view_job, can_modify_job, LinkJobLogs
 from jobbrowser.yarn_models import SparkJob
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 _INITIALIZED = False

+ 0 - 1
apps/jobsub/src/jobsub/models.py

@@ -19,7 +19,6 @@ from builtins import str
 import logging
 
 from django.core import urlresolvers
-from django.contrib.auth.models import User
 from django.db import models
 from django.utils.translation import ugettext_lazy as _
 

+ 7 - 1
apps/jobsub/src/jobsub/tests.py

@@ -22,10 +22,10 @@ import json
 import time
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_raises
-from django.contrib.auth.models import User
 from django.urls import reverse
 from nose.plugins.skip import SkipTest
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_to_group
 from desktop.models import Document
@@ -33,9 +33,15 @@ from desktop.models import Document
 from liboozie.oozie_api_tests import OozieServerProvider
 from oozie.models import Workflow, Node, Start, Kill, End, Link
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 
+
 class TestJobsubWithHadoop(OozieServerProvider):
 
   def setUp(self):

+ 6 - 1
apps/metastore/src/metastore/tests.py

@@ -28,9 +28,9 @@ from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false
 
 from django.utils.encoding import smart_str
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client, assert_equal_mod_whitespace
 from desktop.lib.test_utils import add_permission, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
@@ -44,6 +44,11 @@ from beeswax.models import QueryHistory
 from beeswax.server import dbms
 from beeswax.test_base import BeeswaxSampleProvider
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
apps/oozie/src/oozie/models2_tests.py

@@ -22,14 +22,13 @@ import json
 import logging
 import re
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.db.models import Q
 
 from nose.plugins.attrib import attr
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 
-from desktop.conf import USE_DEFAULT_CONFIGURATION, USE_NEW_EDITOR
+from desktop.conf import USE_DEFAULT_CONFIGURATION, USE_NEW_EDITOR, ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_permission, add_to_group, grant_access, remove_from_group
 from desktop.models import DefaultConfiguration, Document, Document2
@@ -42,6 +41,11 @@ from oozie.models2 import Node, Workflow, WorkflowConfiguration, find_dollar_var
     _create_graph_adjaceny_list, _get_hierarchy_from_adj_list, WorkflowBuilder, WorkflowDepthReached
 from oozie.tests import OozieMockBase, save_temp_workflow, MockOozieApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 7 - 1
apps/oozie/src/oozie/tests.py

@@ -37,11 +37,11 @@ from itertools import chain
 from nose.plugins.skip import SkipTest
 from nose.plugins.attrib import attr
 from nose.tools import raises, assert_true, assert_false, assert_equal, assert_not_equal
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_permission, add_to_group, reformat_json, reformat_xml
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Document, Document2
 import desktop.views as views
 
@@ -68,6 +68,12 @@ if sys.version_info[0] > 2:
 else:
   from cStringIO import StringIO as string_io
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
 

+ 6 - 1
apps/pig/src/pig/tests.py

@@ -20,12 +20,12 @@ from builtins import object
 import json
 import time
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.models import SAMPLE_USER_ID
@@ -39,6 +39,11 @@ from oozie.tests import OozieBase
 from pig.models import create_or_update_script, PigScript
 from pig.api import OozieApi, get
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestPigBase(object):
   SCRIPT_ATTRS = {

+ 6 - 1
apps/search/src/search/tests.py

@@ -19,11 +19,11 @@
 from builtins import object
 import json
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.rest import resource
@@ -33,6 +33,11 @@ from dashboard.facet_builder import _round_number_range
 from dashboard.models import Collection2
 from dashboard.controller import DashboardController
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 QUERY = {'qs': [{'q': ''}], 'fqs': [], 'start': 0}
 

+ 6 - 1
apps/security/src/security/tests.py

@@ -19,16 +19,21 @@
 from builtins import object
 from nose.tools import assert_true, assert_equal, assert_false
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 
 from useradmin.models import HuePermission, GroupPermission
 
 from security.api.hive import _to_sentry_privilege
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 class TestSecurity(object):
 

+ 7 - 2
apps/sqoop/src/sqoop/tests.py

@@ -17,11 +17,11 @@
 import logging
 import json
 
-from django.contrib.auth.models import User
 from nose.tools import assert_true, assert_equal
 from nose.plugins.skip import SkipTest
-
 from django.urls import reverse
+
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 
@@ -29,6 +29,11 @@ from sqoop.client.link import Link
 from sqoop.client.job import Job
 from sqoop.test_base import SqoopServerProvider
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
apps/useradmin/src/useradmin/api.py

@@ -17,12 +17,16 @@
 
 import logging
 
-from django.contrib.auth.models import User, Group
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.decorators import check_superuser_permission
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.i18n import smart_unicode
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 7
apps/useradmin/src/useradmin/forms.py

@@ -21,23 +21,22 @@ import re
 
 import django.contrib.auth.forms
 from django import forms
-from django.contrib.auth.models import User, Group
 from django.forms import ValidationError
 from django.forms.utils import ErrorList
 from django.utils.translation import get_language, ugettext as _, ugettext_lazy as _t
 
 from desktop import conf as desktop_conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import get_username_re_rule, get_groupname_re_rule
 from desktop.settings import LANGUAGES
 
 from useradmin.hue_password_policy import hue_get_password_validators
-from useradmin.models import GroupPermission, HuePermission
-from useradmin.models import get_default_user_group
+from useradmin.models import GroupPermission, HuePermission, get_default_user_group
 
-if desktop_conf.ENABLE_ORGANIZATIONS.get():
-  from useradmin.models2 import OrganizationUser as HueUser
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
 else:
-  from desktop.models import HueUser
+  from django.contrib.auth.models import User, Group
 
 
 LOG = logging.getLogger(__name__)
@@ -120,7 +119,7 @@ class UserChangeForm(django.contrib.auth.forms.UserChangeForm):
                                       required=False)
 
   class Meta(django.contrib.auth.forms.UserChangeForm.Meta):
-    model =  HueUser
+    model =  User
     fields = ["username", "first_name", "last_name", "email", "ensure_home_directory"]
 
   def __init__(self, *args, **kwargs):

+ 10 - 6
apps/useradmin/src/useradmin/ldap_access.py

@@ -25,11 +25,15 @@ import ldap.filter
 import logging
 import re
 
-from django.contrib.auth.models import User
-
 import desktop.conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.python_util import CaseInsensitiveDict
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 
@@ -52,7 +56,7 @@ def get_connection_from_server(server=None):
     ldap_config = ldap_servers[server]
   else:
     ldap_config = desktop.conf.LDAP
-     
+
   return get_connection(ldap_config)
 
 def get_connection(ldap_config):
@@ -211,7 +215,7 @@ class LdapConnection(object):
     """
     :param result_data: List of dictionaries that have ldap attributes and their associated values. Generally the result list from an ldapsearch request.
     :param user_name_attr: The ldap attribute that is returned by the server to map to ``username`` in the return dictionary.
-    
+
     :returns list of dictionaries that take on the following form: {
       'dn': <distinguished name of entry>,
       'username': <ldap attribute associated with user_name_attr>
@@ -310,7 +314,7 @@ class LdapConnection(object):
     :param user_name_attr: The ldap attribute that is returned by the server to map to ``username`` in the return dictionary.
     :param find_by_dn: Search by distinguished name.
     :param scope: ldapsearch scope.
-    
+
     :returns: List of dictionaries that take on the following form: {
       'dn': <distinguished name of entry>,
       'username': <ldap attribute associated with user_name_attr>
@@ -372,7 +376,7 @@ class LdapConnection(object):
     :param group_name_attr: The ldap attribute that is returned by the server to map to ``name`` in the return dictionary.
     :param find_by_dn: Search by distinguished name.
     :param scope: ldapsearch scope.
-    
+
     :returns: List of dictionaries that take on the following form: {
       'dn': <distinguished name of entry>,
       'name': <ldap attribute associated with group_name_attr>

+ 6 - 2
apps/useradmin/src/useradmin/middleware.py

@@ -24,20 +24,24 @@ import logging
 from datetime import datetime
 
 from django.contrib import messages
-from django.contrib.auth.models import User
 from django.contrib.sessions.models import Session
 from django.db import DatabaseError
 from django.db.models import Q
 from django.utils.translation import ugettext as _
 
 from desktop.auth.views import dt_logout
-from desktop.conf import AUTH, LDAP, SESSION
+from desktop.conf import AUTH, LDAP, SESSION, ENABLE_ORGANIZATIONS
 
 from useradmin.models import UserProfile, get_profile
 from useradmin.views import import_ldap_users
 
 from useradmin import ldap_access
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 1
apps/useradmin/src/useradmin/test_ldap.py

@@ -24,10 +24,10 @@ from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false
 
 import desktop.conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import grant_access
 from desktop.lib.django_test_util import make_logged_in_client
 from django.conf import settings
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 
 from useradmin.models import LdapGroup, UserProfile
@@ -41,6 +41,11 @@ from useradmin.views import sync_ldap_users, sync_ldap_groups, import_ldap_users
 from useradmin import ldap_access
 from useradmin.tests import BaseUserAdminTests, LdapTestConnection, reset_all_groups, reset_all_users
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 def get_multi_ldap_config():
   return {'multi_ldap_conf': {

+ 7 - 1
apps/useradmin/src/useradmin/test_ldap_deprecated.py

@@ -24,10 +24,10 @@ from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false
 
 import desktop.conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import grant_access
 from desktop.lib.django_test_util import make_logged_in_client
 from django.conf import settings
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 
 from useradmin.models import LdapGroup, UserProfile, get_profile
@@ -40,6 +40,12 @@ from useradmin.views import sync_ldap_users, sync_ldap_groups, import_ldap_users
 from useradmin import ldap_access
 from useradmin.tests import BaseUserAdminTests, LdapTestConnection, reset_all_groups, reset_all_users
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
+
 class TestUserAdminLdapDeprecated(BaseUserAdminTests):
   def test_useradmin_ldap_user_group_membership_sync(self):
     settings.MIDDLEWARE_CLASSES.append('useradmin.middleware.LdapSynchronizationMiddleware')

+ 6 - 2
apps/useradmin/src/useradmin/tests.py

@@ -30,7 +30,6 @@ import urllib.request, urllib.parse, urllib.error
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false, assert_not_equal
 from datetime import datetime
-from django.contrib.auth.models import User, Group
 from django.contrib.sessions.models import Session
 from django.db.models import Q
 from django.utils.encoding import smart_unicode
@@ -41,7 +40,7 @@ import desktop.conf
 
 from desktop import appmanager
 from desktop.auth.backend import is_admin
-from desktop.conf import APP_BLACKLIST
+from desktop.conf import APP_BLACKLIST ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.views import home
@@ -56,6 +55,11 @@ from useradmin.models import HuePermission, GroupPermission, UserProfile
 from useradmin.models import get_profile, get_default_user_group
 from useradmin.hue_password_policy import reset_password_policy
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 def reset_all_users():
   """Reset to a clean state by deleting all users"""

+ 6 - 1
apps/useradmin/src/useradmin/tests_api.py

@@ -20,10 +20,15 @@ from builtins import object
 import json
 
 from nose.tools import assert_equal, assert_false, assert_true
-from django.contrib.auth.models import User, Group
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 class TestUseradminApi(object):
 

+ 10 - 8
apps/useradmin/src/useradmin/views.py

@@ -24,15 +24,12 @@ import threading
 import subprocess
 import json
 
+
 from axes.decorators import FAILURE_LIMIT, LOCK_OUT_AT_FAILURE
 from axes.models import AccessAttempt
 from axes.utils import reset
-
 import ldap
-from useradmin import ldap_access
-from useradmin.ldap_access import LdapBindException, LdapSearchException
 
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 from django.forms import ValidationError
 from django.forms.utils import ErrorList
@@ -42,21 +39,26 @@ from django.utils.encoding import smart_str
 from django.utils.translation import get_language, ugettext as _
 
 import desktop.conf
-from desktop.conf import LDAP
+from desktop.auth.backend import is_admin
+from desktop.conf import LDAP, ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import JsonResponse, render
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.models import _get_apps
 from desktop.views import antixss
-
 from hadoop.fs.exceptions import WebHdfsException
+
+from useradmin import ldap_access
+from useradmin.ldap_access import LdapBindException, LdapSearchException
 from useradmin.models import HuePermission, UserProfile, LdapGroup
 from useradmin.models import get_profile, get_default_user_group
 from useradmin.forms import SyncLdapUsersGroupsForm, AddLdapGroupsForm, AddLdapUsersForm,\
   PermissionsEditForm, GroupEditForm, SuperUserChangeForm, UserChangeForm, validate_username, validate_first_name, \
   validate_last_name, PasswordChangeForm
 
-from desktop.auth.backend import is_admin
-
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/core/src/desktop/api2_tests.py

@@ -22,12 +22,16 @@ import re
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.models import Document2
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestApi2(object):
 

+ 6 - 3
desktop/core/src/desktop/api_tests.py

@@ -22,10 +22,8 @@ import json
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 from nose.plugins.skip import SkipTest
 
-from django.contrib.auth.models import User
-
 from desktop.api import massaged_documents_for_json, _get_docs
-from desktop.conf import USE_NEW_EDITOR
+from desktop.conf import USE_NEW_EDITOR, ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.models import DocumentTag , Document
@@ -33,6 +31,11 @@ from desktop.models import DocumentTag , Document
 from pig.models import PigScript
 from useradmin.models import get_default_user_group
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestDocModelTags(object):
 

+ 5 - 1
desktop/core/src/desktop/auth/views.py

@@ -32,7 +32,6 @@ import django.contrib.auth.views
 from django.core import urlresolvers
 from django.core.exceptions import SuspiciousOperation
 from django.contrib.auth import login, get_backends, authenticate
-from django.contrib.auth.models import User
 from django.contrib.sessions.models import Session
 from django.http import HttpResponseRedirect
 from django.utils.translation import ugettext as _
@@ -52,6 +51,11 @@ from desktop.lib.django_util import JsonResponse
 from desktop.log.access import access_log, access_warn, last_access_map
 from desktop.settings import LOAD_BALANCER_COOKIE
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 if sys.version_info[0] > 2:
   import urllib.request, urllib.error

+ 6 - 1
desktop/core/src/desktop/auth/views_test.py

@@ -21,7 +21,6 @@ from nose.tools import assert_true, assert_false, assert_equal
 
 from django_auth_ldap import backend as django_auth_ldap_backend
 from django.conf import settings
-from django.contrib.auth.models import User, Group
 from django.test.client import Client
 
 from hadoop.test_base import PseudoHdfsTestBase
@@ -33,9 +32,15 @@ from useradmin.views import import_ldap_groups
 
 from desktop import conf, middleware
 from desktop.auth import backend
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group
+else:
+  from django.contrib.auth.models import User, Group
+
 
 def get_mocked_config():
   return {'mocked_ldap': {

+ 7 - 2
desktop/core/src/desktop/configuration/tests.py

@@ -21,13 +21,18 @@ import json
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.models import DefaultConfiguration
+
 from useradmin.models import get_default_user_group
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestDefaultConfiguration(object):
 

+ 12 - 7
desktop/core/src/desktop/converter_tests.py

@@ -21,21 +21,26 @@ from datetime import datetime
 import json
 
 from nose.tools import assert_equal, assert_false, assert_true
-from django.contrib.auth.models import User
-
-from desktop.converters import DocumentConverter
-from desktop.lib.django_test_util import make_logged_in_client
-from desktop.lib.test_utils import grant_access
-from desktop.models import Directory, Document, Document2, DocumentPermission, DocumentTag
-from librdbms.design import SQLdesign
 
 from beeswax.models import SavedQuery
 from beeswax.design import hql_query
+from librdbms.design import SQLdesign
 from oozie.models import Link, Workflow
 from oozie.tests import add_node
 from pig.models import create_or_update_script
 from useradmin.models import get_default_user_group
 
+from desktop.conf import ENABLE_ORGANIZATIONS
+from desktop.converters import DocumentConverter
+from desktop.lib.django_test_util import make_logged_in_client
+from desktop.lib.test_utils import grant_access
+from desktop.models import Directory, Document, Document2, DocumentPermission, DocumentTag
+
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestDocumentConverter(object):
 

+ 7 - 2
desktop/core/src/desktop/lib/django_test_util.py

@@ -23,10 +23,15 @@ import re
 import json
 
 import django.test.client
-from django.contrib.auth.models import User, Group
-
 import nose.tools
 
+from desktop.conf import ENABLE_ORGANIZATIONS
+
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User, OrganizationGroup as Group, default_organization
+else:
+  from django.contrib.auth.models import User, Group
+
 
 class Client(django.test.client.Client):
   """

+ 9 - 3
desktop/core/src/desktop/lib/fs/proxyfs.py

@@ -22,20 +22,26 @@ from builtins import object
 import logging
 import sys
 
-from django.contrib.auth.models import User
+from urlparse import urlparse
 
 from desktop.auth.backend import is_admin
-from desktop.conf import DEFAULT_USER
+from desktop.conf import DEFAULT_USER, ENABLE_ORGANIZATIONS
+
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
 
 if sys.version_info[0] > 2:
   from urllib.parse import urlparse as lib_urlparse
 else:
   from urlparse import urlparse as lib_urlparse
 
-LOG = logging.getLogger(__name__)
 
+LOG = logging.getLogger(__name__)
 DEFAULT_USER = DEFAULT_USER.get()
 
+
 class ProxyFS(object):
 
   def __init__(self, filesystems_dict, default_scheme, name='default'):

+ 6 - 2
desktop/core/src/desktop/lib/fs/proxyfs_test.py

@@ -22,13 +22,17 @@ from nose.plugins.attrib import attr
 from nose.tools import assert_raises, assert_false, eq_
 from nose import SkipTest
 
-from django.contrib.auth.models import User
 from desktop.auth.backend import rewrite_user
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.fs import ProxyFS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_permission, remove_from_group
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 def test_fs_selection():
   make_logged_in_client(username='test', groupname='default', recreate=True, is_superuser=False)

+ 7 - 1
desktop/core/src/desktop/management/commands/convert_documents.py

@@ -20,12 +20,18 @@ import logging
 import sys
 import time
 
-from django.contrib.auth.models import User
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.converters import DocumentConverter
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 class Command(BaseCommand):
 
   def handle(self, *args, **options):

+ 6 - 1
desktop/core/src/desktop/management/commands/create_user_directories.py

@@ -18,12 +18,17 @@
 
 import logging
 
-from django.contrib.auth.models import User
 from django.core.management.base import CommandError, BaseCommand
 from django.utils.translation import ugettext_lazy as _
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Document2
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/core/src/desktop/metrics.py

@@ -24,10 +24,14 @@ import logging
 import multiprocessing
 import threading
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.metrics import global_registry
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/core/src/desktop/middleware.py

@@ -37,7 +37,6 @@ from django.conf import settings
 from django.contrib import messages
 from django.contrib.auth import REDIRECT_FIELD_NAME, BACKEND_SESSION_KEY, authenticate, load_backend, login
 from django.contrib.auth.middleware import RemoteUserMiddleware
-from django.contrib.auth.models import User
 from django.core import exceptions, urlresolvers
 from django.http import HttpResponseNotAllowed, HttpResponseForbidden
 from django.urls import resolve
@@ -50,6 +49,7 @@ from hadoop import cluster
 import desktop.views
 import desktop.conf
 from desktop.auth.backend import is_admin
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.context_processors import get_app_name
 from desktop.lib import apputil, i18n, fsmanager
 from desktop.lib.django_util import JsonResponse, render, render_json
@@ -59,7 +59,11 @@ from desktop.log import get_audit_logger
 from desktop.log.access import access_log, log_page_hit, access_warn
 from desktop import appmanager
 from desktop import metrics
-from desktop.auth.backend import is_admin
+
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
 
 
 LOG = logging.getLogger(__name__)

+ 6 - 1
desktop/core/src/desktop/monkey_patches.py

@@ -17,11 +17,16 @@
 
 import re
 
-from django.contrib.auth.models import User
 from django.core.validators import RegexValidator
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import get_username_re_rule
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 def monkey_patch_username_validator():
   """

+ 9 - 2
desktop/core/src/desktop/tests.py

@@ -37,7 +37,6 @@ from django.db.models import query, CharField, SmallIntegerField
 from django.core.management import call_command
 from django.core.paginator import Paginator
 from django.conf.urls import url
-from django.contrib.auth.models import User
 from django.db import connection
 from django.urls import reverse
 from django.test.client import Client
@@ -60,7 +59,7 @@ import desktop.views as views
 
 from desktop.auth.backend import rewrite_user
 from desktop.appmanager import DESKTOP_APPS
-from desktop.middleware import DJANGO_VIEW_AUTH_WHITELIST
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.conf import validate_path
 from desktop.lib.django_util import TruncatingModel
@@ -69,6 +68,7 @@ from desktop.lib.conf import _configs_from_dir
 from desktop.lib.paths import get_desktop_root
 from desktop.lib.python_util import force_dict_to_strings
 from desktop.lib.test_utils import grant_access
+from desktop.middleware import DJANGO_VIEW_AUTH_WHITELIST
 from desktop.models import Directory, Document, Document2, get_data_link, _version_from_properties, ClusterConfig, HUE_VERSION
 from desktop.redaction import logfilter
 from desktop.redaction.engine import RedactionPolicy, RedactionRule
@@ -80,8 +80,15 @@ if sys.version_info[0] > 2:
 else:
   from cStringIO import StringIO as string_io
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
+
 def test_home():
   c = make_logged_in_client(username="test_home", groupname="test_home", recreate=True, is_superuser=False)
   user = User.objects.get(username="test_home")

+ 6 - 1
desktop/core/src/desktop/tests_doc2.py

@@ -23,10 +23,10 @@ from datetime import datetime
 
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_false, assert_true, assert_not_equal, assert_raises
-from django.contrib.auth.models import User
 from django.core import management
 from django.db.utils import OperationalError
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.converters import DocumentConverter
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.fs import ProxyFS
@@ -39,6 +39,11 @@ from beeswax.design import hql_query
 from useradmin.models import get_default_user_group
 from oozie.models2 import Workflow
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class MockFs(object):
   def __init__(self):

+ 7 - 2
desktop/libs/aws/src/aws/s3/s3fs_test.py

@@ -20,9 +20,9 @@ import os
 import tempfile
 import string
 
-from django.contrib.auth.models import User
 from nose.tools import assert_equal, assert_false, assert_true, assert_raises, eq_
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_to_group
 from desktop.lib.test_utils import grant_access, add_to_group, add_permission, remove_from_group
@@ -32,6 +32,11 @@ from aws.s3.s3fs import S3FileSystem, S3FileSystemException
 from aws.s3.s3test_utils import S3TestBase, generate_id
 from aws.s3.upload import DEFAULT_WRITE_SIZE
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class S3FSTest(S3TestBase):
 
@@ -288,7 +293,7 @@ class S3FSTest(S3TestBase):
   def test_mkdir(self):
     dir_path = self.get_test_path('test_mkdir')
     assert_false(self.fs.exists(dir_path))
-    
+
     self.fs.mkdir(dir_path)
     assert_true(self.fs.exists(dir_path))
 

+ 6 - 1
desktop/libs/dashboard/src/dashboard/tests.py

@@ -19,11 +19,11 @@
 from builtins import object
 import json
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.rest import resource
@@ -33,6 +33,11 @@ from dashboard.facet_builder import _round_number_range
 from dashboard.models import Collection2, augment_response
 from dashboard.controller import DashboardController
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 QUERY = {'qs': [{'q': ''}], 'fqs': [], 'start': 0}
 

+ 6 - 1
desktop/libs/indexer/src/indexer/indexers/envelope.py

@@ -18,10 +18,10 @@ from builtins import object
 import logging
 import os
 
-from django.contrib.auth.models import User
 from django.core.urlresolvers import reverse
 from django.utils.translation import ugettext as _
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.exceptions_renderable import PopupException
 from hadoop.fs.hadoopfs import Hdfs
 from notebook.models import make_notebook
@@ -29,6 +29,11 @@ from notebook.models import make_notebook
 from indexer.conf import CONFIG_JARS_LIBS_PATH, config_morphline_path
 from libzookeeper.conf import zkensemble
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 42 - 36
desktop/libs/indexer/src/indexer/indexers/envelope_tests.py

@@ -17,13 +17,19 @@
 # limitations under the License.
 
 from builtins import object
-from django.contrib.auth.models import User
 
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 
+from desktop.conf import ENABLE_ORGANIZATIONS
+
 from indexer.indexers.envelope import EnvelopeIndexer
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestEnvelope(object):
 
@@ -34,7 +40,7 @@ class TestEnvelope(object):
   def test_generate_from_kafka_to_file_csv(self):
     properties = {
       'app_name': 'Ingest',
-  
+
       'inputFormat': 'stream',
       'streamSelection': 'kafka',
       'brokers': 'broker:9092',
@@ -43,14 +49,14 @@ class TestEnvelope(object):
       'kafkaFieldDelimiter': ',',
       'kafkaFieldNames': 'id,name',
       'kafkaFieldTypes': 'int,string',
-  
+
       'ouputFormat': 'file',
       'input_path': '/tmp/output',
       'format': 'csv'
     }
-  
+
     config = EnvelopeIndexer(username='test').generate_config(properties)
-  
+
     assert_true('''steps {
       inputdata {
           input {
@@ -68,18 +74,18 @@ class TestEnvelope(object):
                       enabled = true
                       milliseconds = 60000
                   }
-  
+
           }
       }
-  
+
       outputdata {
           dependencies = [inputdata]
-  
+
           deriver {
             type = sql
             query.literal = """SELECT * from inputdata"""
           }
-  
+
           planner = {
             type = overwrite
           }
@@ -97,22 +103,22 @@ class TestEnvelope(object):
   def test_generate_from_stream_sfdc_to_hive_table(self):
     properties = {
       'app_name': 'Ingest',
-  
-      'inputFormat': 'stream',    
+
+      'inputFormat': 'stream',
       'streamSelection': 'sfdc',
       'streamUsername': 'test',
       'streamPassword': 'test',
       'streamToken': 'token',
       'streamEndpointUrl': 'http://sfdc/api',
       'streamObject': 'Opportunities',
-  
+
       'ouputFormat': 'table',
       'output_table': 'sfdc',
       'format': 'text'
     }
-  
+
     config = EnvelopeIndexer(username='test').generate_config(properties)
-  
+
     assert_true('''steps {
       inputdata {
           input {
@@ -127,18 +133,18 @@ class TestEnvelope(object):
                   auth-endpoint = "http://sfdc/api"
                 }
               }
-    
+
           }
       }
-  
+
       outputdata {
           dependencies = [inputdata]
-  
+
           deriver {
             type = sql
             query.literal = """SELECT * from inputdata"""
           }
-  
+
             planner {
                 type = append
             }
@@ -148,12 +154,12 @@ class TestEnvelope(object):
             }
       }
   }''' in  config, config)
-  
+
 
   def test_generate_from_stream_kafka_to_solr_index(self):
     properties = {
       'app_name': 'Ingest',
-  
+
       'inputFormat': 'stream',
       'streamSelection': 'kafka',
       'brokers': 'broker:9092',
@@ -162,14 +168,14 @@ class TestEnvelope(object):
       'kafkaFieldDelimiter': ',',
       'kafkaFieldNames': 'id,name',
       'kafkaFieldTypes': 'int,string',
-  
+
       'ouputFormat': 'index',
       'connection': 'http://hue.com:8983/solr/',
       'collectionName': 'traffic'
     }
-  
+
     config = EnvelopeIndexer(username='test').generate_config(properties)
-  
+
     assert_true('''steps {
       inputdata {
           input {
@@ -187,18 +193,18 @@ class TestEnvelope(object):
                       enabled = true
                       milliseconds = 60000
                   }
-  
+
           }
       }
-  
+
       outputdata {
           dependencies = [inputdata]
-  
+
           deriver {
             type = sql
             query.literal = """SELECT * from inputdata"""
           }
-  
+
           planner {
               type = upstert
           }
@@ -214,39 +220,39 @@ class TestEnvelope(object):
   def test_generate_from_file_to_kafka(self):
     properties = {
       'app_name': 'Ingest',
-  
+
       'inputFormat': 'file',
       'input_path': '/tmp/output',
       'format': 'csv',
-  
+
       'ouputFormat': 'stream',
       'streamSelection': 'kafka',
       'brokers': 'broker:9092',
       'topics': 'kafkaTopic',
       'kafkaFieldType': 'delimited',
     }
-  
+
     config = EnvelopeIndexer(username='test').generate_config(properties)
-  
+
     assert_true('''steps {
       inputdata {
           input {
               type = filesystem
           path = /tmp/output
           format = csv
-        
+
           }
       }
-  
+
       outputdata {
           dependencies = [inputdata]
-  
+
           deriver {
             type = sql
             query.literal = """SELECT * from inputdata"""
           }
-  
-          
+
+
           planner {
               type = append
           }

+ 6 - 1
desktop/libs/indexer/src/indexer/indexers/flume.py

@@ -18,16 +18,21 @@ from builtins import object
 import logging
 import os
 
-from django.contrib.auth.models import User
 from django.core.urlresolvers import reverse
 from django.utils.translation import ugettext as _
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.exceptions_renderable import PopupException
 from libzookeeper.conf import zkensemble
 from metadata.manager_client import ManagerApi
 
 from indexer.conf import config_morphline_path
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 8 - 3
desktop/libs/indexer/src/indexer/indexers/flume_tests.py

@@ -16,13 +16,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from django.contrib.auth.models import User
-
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 
+from desktop.conf import ENABLE_ORGANIZATIONS
+
 from indexer.indexers.flume import FlumeIndexer
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 def test_generate_from_directory_to_solr_index():
   raise SkipTest
@@ -154,7 +159,7 @@ morphlines : [
     ,
     configs[0][1].strip() # 'agent_morphlines_conf_file'
   )
-  
+
   assert_equal(
     ('agent_config_file', 'tier1.sources = source1\n  tier1.channels = channel1\n  tier1.sinks = sink1\n\n\n  tier1.channels.channel1.type = memory\n  tier1.channels.channel1.capacity = 10000\n  tier1.channels.channel1.transactionCapacity = 1000\n\n  \n  tier1.sinks.sink1.type          = org.apache.flume.sink.solr.morphline.MorphlineSolrSink\n  tier1.sinks.sink1.morphlineFile = morphlines.conf\n  tier1.sinks.sink1.morphlineId = hue_accesslogs_no_geo\n  tier1.sinks.sink1.channel       = channel1')
     ,

+ 7 - 3
desktop/libs/indexer/src/indexer/indexers/morphline.py

@@ -20,22 +20,26 @@ import os
 
 from collections import deque
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.utils.translation import ugettext as _
 from mako.lookup import TemplateLookup
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Document2
 from notebook.connectors.base import get_api
 from notebook.models import Notebook, make_notebook
 
-from indexer.conf import CONFIG_INDEXING_TEMPLATES_PATH
-from indexer.conf import CONFIG_INDEXER_LIBS_PATH
+from indexer.conf import CONFIG_INDEXING_TEMPLATES_PATH, CONFIG_INDEXER_LIBS_PATH
 from indexer.fields import get_field_type
 from indexer.file_format import get_file_format_instance, get_file_format_class
 from indexer.indexers.morphline_operations import get_checked_args
 from indexer.solr_client import SolrClient
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 7 - 2
desktop/libs/indexer/src/indexer/indexers/morphline_tests.py

@@ -28,8 +28,7 @@ from nose.tools import assert_equal, assert_true
 from nose.plugins.attrib import attr
 from nose.plugins.skip import SkipTest
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_to_group
 from hadoop.pseudo_hdfs4 import is_live_cluster, shared_cluster
@@ -48,6 +47,12 @@ if sys.version_info[0] > 2:
 else:
   from StringIO import StringIO as string_io
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
 

+ 7 - 2
desktop/libs/indexer/src/indexer/indexers/rdbms.py

@@ -19,10 +19,10 @@ import json
 import logging
 import uuid
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.utils.translation import ugettext as _
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.i18n import smart_str
 from librdbms.conf import DATABASES, get_database_password, get_server_choices, get_connector_name
@@ -32,6 +32,11 @@ from notebook.conf import get_ordered_interpreters
 from notebook.connectors.base import get_api
 from notebook.models import make_notebook
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 
@@ -271,4 +276,4 @@ def _splitby_column_check(statement, destination_splitby_column):
       'statement': statement,
       'destinationSplitbyColumn': destination_splitby_column
     }
-  return statement
+  return statement

+ 7 - 1
desktop/libs/indexer/src/indexer/indexers/sql.py

@@ -22,10 +22,10 @@ import sys
 import urllib.request, urllib.error
 import uuid
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.utils.translation import ugettext as _
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib import django_mako
 from desktop.lib.exceptions_renderable import PopupException
 
@@ -38,6 +38,12 @@ else:
   from urllib import unquote as urllib_unquote
   from urlparse import urlparse
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
 

+ 6 - 1
desktop/libs/indexer/src/indexer/indexers/sql_tests.py

@@ -19,14 +19,19 @@
 from builtins import object
 import json
 
-from django.contrib.auth.models import User
 from mock import patch, Mock, MagicMock
 from nose.tools import assert_equal, assert_true
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 
 from indexer.indexers.sql import SQLIndexer
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class MockRequest(object):
   def __init__(self, fs=None, user=None):

+ 8 - 3
desktop/libs/indexer/src/indexer/rdbms_indexer_tests.py

@@ -19,21 +19,26 @@
 from builtins import object
 import logging
 
-from django.contrib.auth.models import User
-
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_false, assert_not_equal, assert_true
 
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 
 from indexer.conf import ENABLE_SQOOP
 from indexer.indexers.rdbms import _get_api
 from librdbms.server import dbms as rdbms
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 
+
 class TestRdbmsIndexer(object):
 
   @classmethod
@@ -62,4 +67,4 @@ class TestRdbmsIndexer(object):
     assert_not_equal('', data['rows'], data)
 
 class Bag(dict):
-  pass
+  pass

+ 6 - 1
desktop/libs/indexer/src/indexer/solr_client_tests.py

@@ -18,14 +18,19 @@
 from builtins import object
 from nose.tools import assert_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 
 from indexer.solr_client import SolrClient
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestSolrClient(object):
 

+ 6 - 1
desktop/libs/indexer/src/indexer/tests.py

@@ -21,9 +21,9 @@ import json
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster, get_db_prefix
@@ -33,6 +33,11 @@ from libzookeeper import conf as libzookeeper_conf
 from indexer.conf import get_solr_ensemble
 from indexer.controller import CollectionManagerController
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 def test_get_ensemble():
 

+ 15 - 11
desktop/libs/liboauth/src/liboauth/backend.py

@@ -26,16 +26,20 @@ import cgi
 import logging
 import sys
 
-from django.contrib.auth.models import User
 from django.utils.translation import ugettext as _
 
 from desktop.auth.backend import force_username_case, DesktopBackendBase
-from desktop.conf import AUTH
+from desktop.conf import AUTH, ENABLE_ORGANIZATIONS
 from useradmin.models import get_profile, get_default_user_group, UserProfile
 
 import liboauth.conf
 import liboauth.metrics
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 try:
   import oauth2 as oauth
 except:
@@ -69,7 +73,7 @@ class OAuthBackend(DesktopBackendBase):
         is_super = True
       else:
         is_super = False
-    
+
       # Could save oauth_token detail in the user profile here
       user = find_or_create_user(username, password)
 
@@ -86,21 +90,21 @@ class OAuthBackend(DesktopBackendBase):
 
     return user
 
-    
+
   @classmethod
   def manages_passwords_externally(cls):
-    return True 
+    return True
 
   @classmethod
   def is_first_login_ever(cls):
     """ Return true if no external user has ever logged in to Desktop yet. """
     return not UserProfile.objects.filter(creation_method=UserProfile.CreationMethod.EXTERNAL.name).exists()
-  
+
 
   @classmethod
   def handleAuthenticationRequest(cls, request):
     assert oauth is not None
- 
+
     if 'oauth_verifier' in request.GET:
         social = 'twitter'
         consumer_key=liboauth.conf.CONSUMER_KEY_TWITTER.get()
@@ -135,14 +139,14 @@ class OAuthBackend(DesktopBackendBase):
             consumer_key=liboauth.conf.CONSUMER_KEY_GOOGLE.get()
             consumer_secret=liboauth.conf.CONSUMER_SECRET_GOOGLE.get()
             access_token_uri=liboauth.conf.ACCESS_TOKEN_URL_GOOGLE.get()
-            authentication_token_uri=liboauth.conf.AUTHORIZE_URL_GOOGLE.get()        
+            authentication_token_uri=liboauth.conf.AUTHORIZE_URL_GOOGLE.get()
 
         elif social == 'facebook':
             consumer_key=liboauth.conf.CONSUMER_KEY_FACEBOOK.get()
             consumer_secret=liboauth.conf.CONSUMER_SECRET_FACEBOOK.get()
             access_token_uri=liboauth.conf.ACCESS_TOKEN_URL_FACEBOOK.get()
             authentication_token_uri=liboauth.conf.AUTHORIZE_URL_FACEBOOK.get()
-        
+
         elif social == 'linkedin':
             consumer_key=liboauth.conf.CONSUMER_KEY_LINKEDIN.get()
             consumer_secret=liboauth.conf.CONSUMER_SECRET_LINKEDIN.get()
@@ -192,7 +196,7 @@ class OAuthBackend(DesktopBackendBase):
                 raise Exception(_("Invalid response from OAuth provider: %s") % resp)
             username = (json.loads(content))['emailAddress']
             access_token = dict(screen_name=map_username(username), oauth_token_secret=access_tok)
-  
+
 
     return access_token, nexturl
 
@@ -203,7 +207,7 @@ class OAuthBackend(DesktopBackendBase):
 
     redirect_uri = get_redirect_uri(request)
     response_type = "code"
- 
+
     social = request.GET['social']
     state = social + "," + request.GET.get('next', '/')
 

+ 9 - 4
desktop/libs/liboauth/src/liboauth/views.py

@@ -21,7 +21,7 @@ try:
   import oauth2 as oauth
 except:
   oauth = None
- 
+
 import logging
 import urllib.request, urllib.parse, urllib.error
 import httplib2
@@ -30,7 +30,6 @@ import django.contrib.auth.views
 from django.core import urlresolvers
 from django.core.exceptions import SuspiciousOperation
 from django.contrib.auth import login, get_backends, authenticate
-from django.contrib.auth.models import User
 from django.contrib.sessions.models import Session
 from django.http import HttpResponseRedirect
 from django.utils.translation import ugettext as _
@@ -39,6 +38,7 @@ from useradmin.views import ensure_home_directory
 
 from desktop.auth.backend import AllowFirstUserDjangoBackend
 from desktop.auth.forms import UserCreationForm, AuthenticationForm
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import render
 from desktop.lib.django_util import login_notrequired
 from desktop.log.access import access_warn, last_access_map
@@ -46,6 +46,11 @@ from desktop.log.access import access_warn, last_access_map
 import liboauth.conf
 from liboauth.backend import OAuthBackend
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 @login_notrequired
 def show_login_page(request, login_errors=False):
@@ -77,10 +82,10 @@ def oauth_login(request):
 
   return HttpResponseRedirect(url)
 
-  
+
 @login_notrequired
 def oauth_authenticated(request):
-   
+
   access_token, next = OAuthBackend.handleAuthenticationRequest(request)
   if access_token == "":
       return show_login_page(request, True)

+ 6 - 1
desktop/libs/liboozie/src/liboozie/submittion2_tests.py

@@ -19,7 +19,6 @@ from __future__ import print_function
 from builtins import object
 import logging
 
-from django.contrib.auth.models import User
 from nose.plugins.attrib import attr
 from nose.tools import assert_equal, assert_true, assert_not_equal, assert_raises
 
@@ -28,6 +27,7 @@ import beeswax
 from hadoop import cluster, pseudo_hdfs4
 from hadoop.conf import HDFS_CLUSTERS, MR_CLUSTERS, YARN_CLUSTERS
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import clear_sys_caches
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.exceptions_renderable import PopupException
@@ -40,6 +40,11 @@ from liboozie.credentials import Credentials
 from liboozie.credentials_tests import TestCredentials
 from liboozie.submission2 import Submission
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 1
desktop/libs/liboozie/src/liboozie/submittion_tests.py

@@ -18,13 +18,13 @@
 from builtins import object
 import logging
 
-from django.contrib.auth.models import User
 from nose.plugins.attrib import attr
 from nose.tools import assert_equal, assert_true, assert_not_equal
 
 from hadoop import cluster, pseudo_hdfs4
 from hadoop.conf import HDFS_CLUSTERS, MR_CLUSTERS, YARN_CLUSTERS
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import clear_sys_caches
 from desktop.lib.django_test_util import make_logged_in_client
 from oozie.tests import OozieMockBase
@@ -32,6 +32,11 @@ from useradmin.views import ensure_home_directory
 
 from liboozie.submittion import Submission
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 1
desktop/libs/liboozie/src/liboozie/tests.py

@@ -17,10 +17,10 @@
 
 import logging
 
-from django.contrib.auth.models import User
 from nose.tools import assert_equal
 
 import desktop.conf as desktop_conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import reformat_xml
 
 from liboozie import conf
@@ -28,6 +28,11 @@ from liboozie.types import WorkflowAction, Coordinator
 from liboozie.utils import config_gen
 from oozie.tests import MockOozieApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 1
desktop/libs/libopenid/src/libopenid/backend.py

@@ -23,14 +23,19 @@ import logging
 import sys
 
 from django.contrib.auth import logout as auth_logout
-from django.contrib.auth.models import User
 from django_openid_auth.auth import OpenIDBackend as _OpenIDBackend
 
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from useradmin.models import get_profile, get_default_user_group, UserProfile
 
 from libopenid import metrics
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 7 - 2
desktop/libs/libsaml/src/libsaml/backend.py

@@ -23,18 +23,23 @@ from __future__ import absolute_import
 import logging
 
 from django.contrib.auth import logout as auth_logout
-from django.contrib.auth.models import User
 from djangosaml2.backends import Saml2Backend as _Saml2Backend
 from djangosaml2.views import logout as saml_logout
 
 from desktop.auth.backend import force_username_case, rewrite_user
-from desktop.conf import AUTH
+from desktop.conf import AUTH, ENABLE_ORGANIZATIONS
 
 from libsaml import conf
 from libsaml import metrics
 
 from useradmin.models import get_profile, get_default_user_group, UserProfile
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
 

+ 6 - 1
desktop/libs/libsentry/src/libsentry/test_privilege_checker.py

@@ -19,14 +19,19 @@
 from builtins import object
 import pickle
 
-from django.contrib.auth.models import User
 from nose.tools import assert_equal, assert_false, assert_true
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 
 from libsentry.privilege_checker import PrivilegeChecker
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class MockSentryApiV1(object):
 

+ 6 - 2
desktop/libs/libsentry/src/libsentry/tests.py

@@ -24,8 +24,7 @@ import tempfile
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false, assert_not_equal, assert_raises
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.exceptions import StructuredThriftTransportException
 from desktop.lib.exceptions_renderable import PopupException
@@ -41,6 +40,11 @@ from libsentry.client import SentryClient
 from libsentry.sentry_ha import get_next_available_server
 from libsentry.sentry_site import get_sentry_server
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 def create_mock_client_fn(client_class, username, server, component=None):
   class MockSentryClient(object):

+ 6 - 1
desktop/libs/libsolr/src/libsolr/tests.py

@@ -22,16 +22,21 @@ import json
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from dashboard.models import Collection2
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
 
 from libsolr.api import SolrApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/libs/libzookeeper/src/libzookeeper/tests.py

@@ -23,8 +23,7 @@ import tempfile
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
@@ -32,6 +31,11 @@ from hadoop.pseudo_hdfs4 import is_live_cluster
 from libzookeeper.models import ZookeeperClient
 from libzookeeper.conf import zkensemble, ENSEMBLE
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class UnitTests(object):
 

+ 6 - 1
desktop/libs/metadata/src/metadata/catalog/navigator_client_tests.py

@@ -21,9 +21,9 @@ from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal
 
 from django.core.cache import cache
-from django.contrib.auth.models import User
 
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from libsentry.privilege_checker import SENTRY_PRIVILEGE_CACHE_KEY, PrivilegeChecker
@@ -33,6 +33,11 @@ from metadata.conf import NAVIGATOR
 from metadata.metadata_sites import get_navigator_hue_server_name
 from metadata.catalog.navigator_client import NavigatorApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 5 - 1
desktop/libs/metadata/src/metadata/catalog_tests.py

@@ -22,10 +22,10 @@ import json
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
@@ -36,6 +36,10 @@ from metadata.catalog_api import _augment_highlighting
 from metadata.catalog.navigator_client import NavigatorApi
 from metadata.catalog.atlas_client import AtlasApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
 
 
 LOG = logging.getLogger(__name__)

+ 6 - 2
desktop/libs/metadata/src/metadata/optimizer_api_tests.py

@@ -19,14 +19,18 @@ import logging
 
 from nose.tools import assert_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
-
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 
 from metadata.optimizer_api import _convert_queries
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/libs/metadata/src/metadata/optimizer_client_tests.py

@@ -21,9 +21,8 @@ import time
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
-
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 
@@ -31,6 +30,11 @@ from metadata.optimizer_client import OptimizerApi
 from metadata.conf import OPTIMIZER, has_optimizer
 from desktop.models import uuid_default
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/libs/notebook/src/notebook/connectors/hiveserver2_tests.py

@@ -27,12 +27,11 @@ from mock import patch, Mock
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from TCLIService.ttypes import TStatusCode, TProtocolVersion, TOperationType
 
 from desktop.auth.backend import rewrite_user
-from desktop.conf import has_connectors
+from desktop.conf import has_connectors, ENABLE_ORGANIZATIONS
 from desktop.lib.i18n import smart_str
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
@@ -45,6 +44,11 @@ from notebook.models import make_notebook, Notebook
 from beeswax.server import dbms
 from beeswax.test_base import BeeswaxSampleProvider, get_query_server_config, is_hive_on_spark
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 2
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -22,14 +22,18 @@ import logging
 from mock import patch, Mock, MagicMock
 from nose.tools import assert_equal, assert_not_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
-
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 
 from notebook.connectors.sql_alchemy import SqlAlchemyApi
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 6 - 1
desktop/libs/notebook/src/notebook/management/commands/notebook_setup.py

@@ -19,14 +19,19 @@ import logging
 import os
 import pwd
 
-from django.contrib.auth.models import User
 from django.core import management
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Directory, Document, Document2, Document2Permission, SAMPLE_USER_OWNERS
 from useradmin.models import get_default_user_group, install_sample_user
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 LOG = logging.getLogger(__name__)
 

+ 7 - 2
desktop/libs/notebook/src/notebook/models.py

@@ -29,14 +29,13 @@ import uuid
 
 from datetime import timedelta
 
-from django.contrib.auth.models import User
 from django.contrib.sessions.models import Session
 from django.db.models import Count
 from django.db.models.functions import Trunc
 from django.utils.html import escape
 from django.utils.translation import ugettext as _
 
-from desktop.conf import has_connectors
+from desktop.conf import has_connectors, ENABLE_ORGANIZATIONS
 from desktop.lib.i18n import smart_unicode
 from desktop.lib.paths import SAFE_CHARACTERS_URI
 from desktop.models import Document2
@@ -49,6 +48,12 @@ if sys.version_info[0] > 2:
 else:
   from urllib import quote as urllib_quote
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG = logging.getLogger(__name__)
 
 

+ 7 - 1
desktop/libs/notebook/src/notebook/tasks.py

@@ -32,12 +32,12 @@ from celery import states
 
 from django.core.cache import caches
 from django.core.files.storage import get_storage_class
-from django.contrib.auth.models import User
 from django.db import transaction
 from django.http import FileResponse, HttpRequest
 
 from beeswax import data_export
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.celery import app
 from desktop.conf import TASK_SERVER
 from desktop.lib import export_csvxls
@@ -52,6 +52,12 @@ if sys.version_info[0] > 2:
 else:
   from StringIO import StringIO as string_io
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
+
 LOG_TASK = get_task_logger(__name__)
 LOG = logging.getLogger(__name__)
 STATE_MAP = {

+ 6 - 2
desktop/libs/notebook/src/notebook/tests.py

@@ -24,12 +24,11 @@ from mock import patch, Mock, MagicMock
 from nose.plugins.attrib import attr
 from nose.tools import assert_equal, assert_true, assert_false
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from azure.conf import is_adls_enabled
 
 from desktop import appmanager
-from desktop.conf import APP_BLACKLIST
+from desktop.conf import APP_BLACKLIST, ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_permission
 from desktop.models import Directory, Document, Document2
@@ -43,6 +42,11 @@ from notebook.decorators import api_error_handler
 from notebook.conf import get_ordered_interpreters, INTERPRETERS_SHOWN_ON_WHEEL, INTERPRETERS
 from notebook.models import Analytics
 
+if ENABLE_ORGANIZATIONS.get():
+  from useradmin.models2 import OrganizationUser as User
+else:
+  from django.contrib.auth.models import User
+
 
 class TestNotebookApi(object):