Просмотр исходного кода

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

Romain 6 лет назад
Родитель
Сommit
672df534b9
82 измененных файлов с 594 добавлено и 217 удалено
  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
 from builtins import object
 import json
 import json
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 from nose.tools import assert_true, assert_false, assert_equal
 from nose.tools import assert_true, assert_false, assert_equal
 
 
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Settings
 from desktop.models import Settings
 from oozie.tests import OozieBase
 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):
 class TestAboutBase(object):

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

@@ -20,13 +20,14 @@ import logging
 import json
 import json
 import re
 import re
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.http import Http404
 from django.http import Http404
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 from django.views.decorators.http import require_POST
 from django.views.decorators.http import require_POST
 
 
 from thrift.transport.TTransport import TTransportException
 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.context_processors import get_app_name
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.exceptions import StructuredThriftTransportException
 from desktop.lib.exceptions import StructuredThriftTransportException
@@ -37,7 +38,6 @@ from metastore import parser
 from notebook.models import escape_rows, MockedDjangoRequest, make_notebook
 from notebook.models import escape_rows, MockedDjangoRequest, make_notebook
 
 
 import beeswax.models
 import beeswax.models
-
 from beeswax.data_export import upload
 from beeswax.data_export import upload
 from beeswax.design import HQLdesign
 from beeswax.design import HQLdesign
 from beeswax.conf import USE_GET_LOG_API
 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.models import Session, QueryHistory
 from beeswax.server import dbms
 from beeswax.server import dbms
 from beeswax.server.dbms import expand_exception, get_query_server_config, QueryServerException, QueryServerTimeoutException,\
 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,\
 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.conf import FORCE_HS2_METADATA
 from metastore.views import _get_db, _get_servername
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -23,11 +23,10 @@ import os
 import pwd
 import pwd
 
 
 from django.core.management.base import BaseCommand
 from django.core.management.base import BaseCommand
-from django.contrib.auth.models import User
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
 from desktop.lib.exceptions_renderable import PopupException
 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 desktop.models import Directory, Document, Document2, Document2Permission
 from hadoop import cluster
 from hadoop import cluster
 from notebook.models import import_saved_beeswax_query
 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 import dbms
 from beeswax.server.dbms import get_query_server_config, QueryServerException
 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__)
 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 nose.tools import assert_true, assert_false
 from django.urls import reverse
 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.django_test_util import make_logged_in_client
 from desktop.lib.paths import get_run_root
 from desktop.lib.paths import get_run_root
 from desktop.lib.python_util import find_unused_port
 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
 from hadoop.pseudo_hdfs4 import is_live_cluster, get_db_prefix
 
 
 import beeswax.conf
 import beeswax.conf
-
 from beeswax.server.dbms import get_query_server_config
 from beeswax.server.dbms import get_query_server_config
 from beeswax.server import dbms
 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()
 HIVE_SERVER_TEST_PORT = find_unused_port()
 _INITIALIZED = False
 _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.encoding import smart_str
 from django.utils.html import escape
 from django.utils.html import escape
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.db import transaction
 from django.db import transaction
 
 
@@ -55,7 +54,7 @@ from desktop.conf import \
     AUTH_PASSWORD_SCRIPT as DEFAULT_AUTH_PASSWORD_SCRIPT, \
     AUTH_PASSWORD_SCRIPT as DEFAULT_AUTH_PASSWORD_SCRIPT, \
     LDAP_USERNAME, \
     LDAP_USERNAME, \
     LDAP_PASSWORD, \
     LDAP_PASSWORD, \
-    USE_NEW_EDITOR
+    USE_NEW_EDITOR, ENABLE_ORGANIZATIONS
 from desktop import redaction
 from desktop import redaction
 from desktop.redaction import logfilter
 from desktop.redaction import logfilter
 from desktop.redaction.engine import RedactionPolicy, RedactionRule
 from desktop.redaction.engine import RedactionPolicy, RedactionRule
@@ -98,8 +97,15 @@ if sys.version_info[0] > 2:
 else:
 else:
   from cStringIO import StringIO as string_io
   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__)
 LOG = logging.getLogger(__name__)
 
 
+
 def _list_dir_without_temp_files(fs, target_dir):
 def _list_dir_without_temp_files(fs, target_dir):
   return [f for f in fs.listdir(target_dir) if not f.startswith('.')]
   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 import forms
 from django.core.paginator import Paginator, EmptyPage, InvalidPage
 from django.core.paginator import Paginator, EmptyPage, InvalidPage
 from django.contrib import messages
 from django.contrib import messages
-from django.contrib.auth.models import User
 from django.db.models import Q
 from django.db.models import Q
 from django.http import HttpResponse, QueryDict
 from django.http import HttpResponse, QueryDict
 from django.shortcuts import redirect
 from django.shortcuts import redirect
@@ -35,7 +34,7 @@ from django.utils.translation import ugettext as _
 from django.urls import reverse
 from django.urls import reverse
 
 
 from desktop.appmanager import get_apps_dict
 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.context_processors import get_app_name
 
 
 from desktop.lib.django_util import JsonResponse
 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
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -24,23 +24,30 @@ import sys
 import urllib.request, urllib.error
 import urllib.request, urllib.error
 
 
 from django import forms
 from django import forms
-from django.contrib.auth.models import User, Group
 from django.forms import FileField, CharField, BooleanField, Textarea
 from django.forms import FileField, CharField, BooleanField, Textarea
 from django.forms.formsets import formset_factory, BaseFormSet
 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 aws.s3 import S3A_ROOT, normpath as s3_normpath
 from azure.abfs.__init__ import ABFS_ROOT, normpath as abfs_normpath
 from azure.abfs.__init__ import ABFS_ROOT, normpath as abfs_normpath
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib import i18n
 from desktop.lib import i18n
 from hadoop.fs import normpath
 from hadoop.fs import normpath
+
 from filebrowser.lib import rwx
 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:
 if sys.version_info[0] > 2:
   from urllib.parse import unquote as urllib_unquote
   from urllib.parse import unquote as urllib_unquote
 else:
 else:
   from urllib import unquote as urllib_unquote
   from urllib import unquote as urllib_unquote
 
 
+
 logger = logging.getLogger(__name__)
 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 datetime import datetime
 from gzip import GzipFile
 from gzip import GzipFile
 
 
-from django.contrib.auth.models import User, Group
 from django.core.paginator import EmptyPage, Paginator, Page, InvalidPage
 from django.core.paginator import EmptyPage, Paginator, Page, InvalidPage
 from django.urls import reverse
 from django.urls import reverse
 from django.template.defaultfilters import stringformat, filesizeformat
 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 aws.s3.s3fs import S3FileSystemException, S3ListAllBucketsException
 from desktop import appmanager
 from desktop import appmanager
 from desktop.auth.backend import is_admin
 from desktop.auth.backend import is_admin
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib import i18n
 from desktop.lib import i18n
 from desktop.lib.conf import coerce_bool
 from desktop.lib.conf import coerce_bool
 from desktop.lib.django_util import render, format_preserving_redirect
 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.rwx import filetype, rwx
 from filebrowser.lib import xxd
 from filebrowser.lib import xxd
 from filebrowser.forms import RenameForm, UploadFileForm, UploadArchiveForm, MkDirForm, EditorForm, TouchForm,\
 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:
 if sys.version_info[0] > 2:
   import io
   import io
   from io import StringIO as string_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.s3fs import S3FileSystemException
 from aws.s3.s3test_utils import get_test_bucket
 from aws.s3.s3test_utils import get_test_bucket
+
 from azure.conf import is_abfs_enabled, is_adls_enabled
 from azure.conf import is_abfs_enabled, is_adls_enabled
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.encoding import smart_str
 from django.utils.encoding import smart_str
+from django.utils.translation import ugettext_lazy as _
 from nose.plugins.attrib import attr
 from nose.plugins.attrib import attr
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises,\
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises,\
   assert_greater
   assert_greater
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_test_util import make_logged_in_client
 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.test_utils import grant_access, add_to_group, add_permission, remove_from_group
 from desktop.lib.view_util import location_to_url
 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 filebrowser.views import snappy_installed
 from desktop.conf import is_oozie_enabled
 from desktop.conf import is_oozie_enabled
 
 
-
-
 if sys.version_info[0] > 2:
 if sys.version_info[0] > 2:
   from urllib.parse import unquote as urllib_unquote
   from urllib.parse import unquote as urllib_unquote
 else:
 else:
   from urllib import unquote as urllib_unquote
   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__)
 LOG = logging.getLogger(__name__)
 
 
@@ -982,7 +987,7 @@ alert("XSS")
   def test_compress_hdfs_files(self):
   def test_compress_hdfs_files(self):
     if not is_oozie_enabled():
     if not is_oozie_enabled():
       raise SkipTest
       raise SkipTest
-    
+
     def make_and_test_dir(pre, test_direct):
     def make_and_test_dir(pre, test_direct):
       test_dir = pre + "/" + test_direct
       test_dir = pre + "/" + test_direct
       test_file = test_dir + '/test.txt'
       test_file = test_dir + '/test.txt'
@@ -992,12 +997,12 @@ alert("XSS")
       for i in range(3):
       for i in range(3):
         f = self.cluster.fs.open(test_file + "%s" %i, "w")
         f = self.cluster.fs.open(test_file + "%s" %i, "w")
         f.close()
         f.close()
-        
+
       resp = self.c.post('/filebrowser/compress_files', {'upload_path': pre, 'files[]': [test_direct], 'archive_name': 'test_compress.zip'})
       resp = self.c.post('/filebrowser/compress_files', {'upload_path': pre, 'files[]': [test_direct], 'archive_name': 'test_compress.zip'})
       response = json.loads(resp.content)
       response = json.loads(resp.content)
       assert_equal(0, response['status'], response)
       assert_equal(0, response['status'], response)
       assert_true('handle' in response and response['handle']['id'], response)
       assert_true('handle' in response and response['handle']['id'], response)
-      responseid = '"' + response['handle']['id'] + '"' 
+      responseid = '"' + response['handle']['id'] + '"'
       timeout_time = time() + 25
       timeout_time = time() + 25
       end_time = time()
       end_time = time()
       while timeout_time > end_time:
       while timeout_time > end_time:
@@ -1005,16 +1010,16 @@ alert("XSS")
         response2 = json.loads(resp2.content)
         response2 = json.loads(resp2.content)
         if response2['app']['status'] != 'RUNNING':
         if response2['app']['status'] != 'RUNNING':
           assert_equal(response2['app']['status'] , 'SUCCEEDED', response2)
           assert_equal(response2['app']['status'] , 'SUCCEEDED', response2)
-          break 
+          break
         sleep(3)
         sleep(3)
         end_time = time()
         end_time = time()
       assert_greater(timeout_time, end_time, response)
       assert_greater(timeout_time, end_time, response)
-        
-        
+
+
     ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(True)
     ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(True)
     prefix = self.cluster.fs_prefix + '/test_compress_files'
     prefix = self.cluster.fs_prefix + '/test_compress_files'
     self.cluster.fs.mkdir(prefix)
     self.cluster.fs.mkdir(prefix)
-    
+
     try:
     try:
       make_and_test_dir(prefix, 'testdir')
       make_and_test_dir(prefix, 'testdir')
       make_and_test_dir(prefix, 'test dir1')
       make_and_test_dir(prefix, 'test dir1')
@@ -1023,7 +1028,7 @@ alert("XSS")
     finally:
     finally:
       ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(False)
       ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(False)
       cleanup_tree(self.cluster, prefix)
       cleanup_tree(self.cluster, prefix)
-      
+
 
 
   def test_extract_tgz(self):
   def test_extract_tgz(self):
     ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(True)
     ENABLE_EXTRACT_UPLOADED_ARCHIVE.set_for_testing(True)
@@ -1243,7 +1248,7 @@ class TestS3AccessPermissions(object):
       assert_equal(200, response.status_code)
       assert_equal(200, response.status_code)
     finally:
     finally:
       remove_from_group(self.user.username, 'has_s3')
       remove_from_group(self.user.username, 'has_s3')
-      
+
 class TestABFSAccessPermissions(object):
 class TestABFSAccessPermissions(object):
 
 
   def setUp(self):
   def setUp(self):
@@ -1275,7 +1280,7 @@ class TestABFSAccessPermissions(object):
       assert_equal(200, response.status_code)
       assert_equal(200, response.status_code)
     finally:
     finally:
       remove_from_group(self.user.username, 'has_abfs')
       remove_from_group(self.user.username, 'has_abfs')
-      
+
 class TestADLSAccessPermissions(object):
 class TestADLSAccessPermissions(object):
 
 
   def setUp(self):
   def setUp(self):
@@ -1290,7 +1295,7 @@ class TestADLSAccessPermissions(object):
       raise SkipTest
       raise SkipTest
     response = self.client.get('/filebrowser/view=ADL://')
     response = self.client.get('/filebrowser/view=ADL://')
     assert_equal(500, response.status_code)
     assert_equal(500, response.status_code)
-    
+
     response = self.client.get('/filebrowser/view=ADL://hue_adls_testing')
     response = self.client.get('/filebrowser/view=ADL://hue_adls_testing')
     assert_equal(500, response.status_code)
     assert_equal(500, response.status_code)
 
 
@@ -1321,4 +1326,4 @@ class TestADLSAccessPermissions(object):
       response = self.client.get('/filebrowser/view=ADL://')
       response = self.client.get('/filebrowser/view=ADL://')
       assert_equal(200, response.status_code)
       assert_equal(200, response.status_code)
     finally:
     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 datetime import datetime, timedelta
 
 
-from django.contrib.auth.models import User
 from django.core.management.base import BaseCommand
 from django.core.management.base import BaseCommand
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.paths import get_apps_root
 from desktop.lib.paths import get_apps_root
 
 
 from useradmin.models import install_sample_user
 from useradmin.models import install_sample_user
@@ -31,6 +31,11 @@ from useradmin.models import install_sample_user
 from hbased.ttypes import AlreadyExists
 from hbased.ttypes import AlreadyExists
 from hbase.api import HbaseApi
 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__)
 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 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.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
 
 
@@ -35,6 +34,11 @@ from hadoop.pseudo_hdfs4 import is_live_cluster
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
 from django.urls import reverse
 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():
 def test_security_plain():
   tmpdir = tempfile.mkdtemp()
   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 nose.tools import assert_true, assert_equal, assert_false, assert_raises
 from mock import patch, Mock
 from mock import patch, Mock
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 import desktop.conf as desktop_conf
 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.django_test_util import make_logged_in_client
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.test_utils import add_to_group
 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 import data_export
 from beeswax.design import hql_query
 from beeswax.design import hql_query
-
 from beeswax.data_export import download
 from beeswax.data_export import download
 from beeswax.models import SavedQuery, QueryHistory
 from beeswax.models import SavedQuery, QueryHistory
 from beeswax.server import dbms
 from beeswax.server import dbms
@@ -46,6 +45,11 @@ from beeswax.tests import _make_query
 from impala import conf
 from impala import conf
 from impala.dbms import ImpalaDbms
 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__)
 LOG = logging.getLogger(__name__)
 
 
@@ -516,4 +520,3 @@ class TestImpalaDbms(object):
                  ('order_date', '`default`.`customers`.`orders`'))
                  ('order_date', '`default`.`customers`.`orders`'))
     assert_equal(ImpalaDbms.get_nested_select('default', 'customers', 'orders', 'item/items/item/product_id'),
     assert_equal(ImpalaDbms.get_nested_select('default', 'customers', 'orders', 'item/items/item/product_id'),
                  ('product_id', '`default`.`customers`.`orders`.`items`'))
                  ('product_id', '`default`.`customers`.`orders`.`items`'))
-

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

@@ -24,11 +24,11 @@ import re
 import time
 import time
 import unittest
 import unittest
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_false, assert_equal, assert_raises
 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.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
 from desktop.models import Document
 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.models import can_view_job, can_modify_job, LinkJobLogs
 from jobbrowser.yarn_models import SparkJob
 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__)
 LOG = logging.getLogger(__name__)
 _INITIALIZED = False
 _INITIALIZED = False

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

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

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

@@ -22,10 +22,10 @@ import json
 import time
 import time
 
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_raises
 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 django.urls import reverse
 from nose.plugins.skip import SkipTest
 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.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
 from desktop.models import Document
 from desktop.models import Document
@@ -33,9 +33,15 @@ from desktop.models import Document
 from liboozie.oozie_api_tests import OozieServerProvider
 from liboozie.oozie_api_tests import OozieServerProvider
 from oozie.models import Workflow, Node, Start, Kill, End, Link
 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__)
 LOG = logging.getLogger(__name__)
 
 
+
 class TestJobsubWithHadoop(OozieServerProvider):
 class TestJobsubWithHadoop(OozieServerProvider):
 
 
   def setUp(self):
   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 nose.tools import assert_true, assert_equal, assert_false
 
 
 from django.utils.encoding import smart_str
 from django.utils.encoding import smart_str
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 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.django_test_util import make_logged_in_client, assert_equal_mod_whitespace
 from desktop.lib.test_utils import add_permission, grant_access
 from desktop.lib.test_utils import add_permission, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
 from hadoop.pseudo_hdfs4 import is_live_cluster
@@ -44,6 +44,11 @@ from beeswax.models import QueryHistory
 from beeswax.server import dbms
 from beeswax.server import dbms
 from beeswax.test_base import BeeswaxSampleProvider
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -22,14 +22,13 @@ import json
 import logging
 import logging
 import re
 import re
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.db.models import Q
 from django.db.models import Q
 
 
 from nose.plugins.attrib import attr
 from nose.plugins.attrib import attr
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 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.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.lib.test_utils import add_permission, add_to_group, grant_access, remove_from_group
 from desktop.models import DefaultConfiguration, Document, Document2
 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
     _create_graph_adjaceny_list, _get_hierarchy_from_adj_list, WorkflowBuilder, WorkflowDepthReached
 from oozie.tests import OozieMockBase, save_temp_workflow, MockOozieApi
 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__)
 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.skip import SkipTest
 from nose.plugins.attrib import attr
 from nose.plugins.attrib import attr
 from nose.tools import raises, assert_true, assert_false, assert_equal, assert_not_equal
 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 django.urls import reverse
 
 
 from desktop.lib.django_test_util import make_logged_in_client
 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.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
 from desktop.models import Document, Document2
 import desktop.views as views
 import desktop.views as views
 
 
@@ -68,6 +68,12 @@ if sys.version_info[0] > 2:
 else:
 else:
   from cStringIO import StringIO as string_io
   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__)
 LOG = logging.getLogger(__name__)
 
 
 
 

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

@@ -20,12 +20,12 @@ from builtins import object
 import json
 import json
 import time
 import time
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.models import SAMPLE_USER_ID
 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.models import create_or_update_script, PigScript
 from pig.api import OozieApi, get
 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):
 class TestPigBase(object):
   SCRIPT_ATTRS = {
   SCRIPT_ATTRS = {

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

@@ -19,11 +19,11 @@
 from builtins import object
 from builtins import object
 import json
 import json
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.lib.rest import resource
 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.models import Collection2
 from dashboard.controller import DashboardController
 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}
 QUERY = {'qs': [{'q': ''}], 'fqs': [], 'start': 0}
 
 

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

@@ -19,16 +19,21 @@
 from builtins import object
 from builtins import object
 from nose.tools import assert_true, assert_equal, assert_false
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 
 
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 from django.urls import reverse
 
 
 from useradmin.models import HuePermission, GroupPermission
 from useradmin.models import HuePermission, GroupPermission
 
 
 from security.api.hive import _to_sentry_privilege
 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):
 class TestSecurity(object):
 
 

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

@@ -17,11 +17,11 @@
 import logging
 import logging
 import json
 import json
 
 
-from django.contrib.auth.models import User
 from nose.tools import assert_true, assert_equal
 from nose.tools import assert_true, assert_equal
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
-
 from django.urls import reverse
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 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.client.job import Job
 from sqoop.test_base import SqoopServerProvider
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -17,12 +17,16 @@
 
 
 import logging
 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.decorators import check_superuser_permission
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.i18n import smart_unicode
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

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

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

@@ -25,11 +25,15 @@ import ldap.filter
 import logging
 import logging
 import re
 import re
 
 
-from django.contrib.auth.models import User
-
 import desktop.conf
 import desktop.conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.python_util import CaseInsensitiveDict
 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__)
 LOG = logging.getLogger(__name__)
 
 
@@ -52,7 +56,7 @@ def get_connection_from_server(server=None):
     ldap_config = ldap_servers[server]
     ldap_config = ldap_servers[server]
   else:
   else:
     ldap_config = desktop.conf.LDAP
     ldap_config = desktop.conf.LDAP
-     
+
   return get_connection(ldap_config)
   return get_connection(ldap_config)
 
 
 def 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 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.
     :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: {
     :returns list of dictionaries that take on the following form: {
       'dn': <distinguished name of entry>,
       'dn': <distinguished name of entry>,
       'username': <ldap attribute associated with user_name_attr>
       '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 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 find_by_dn: Search by distinguished name.
     :param scope: ldapsearch scope.
     :param scope: ldapsearch scope.
-    
+
     :returns: List of dictionaries that take on the following form: {
     :returns: List of dictionaries that take on the following form: {
       'dn': <distinguished name of entry>,
       'dn': <distinguished name of entry>,
       'username': <ldap attribute associated with user_name_attr>
       '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 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 find_by_dn: Search by distinguished name.
     :param scope: ldapsearch scope.
     :param scope: ldapsearch scope.
-    
+
     :returns: List of dictionaries that take on the following form: {
     :returns: List of dictionaries that take on the following form: {
       'dn': <distinguished name of entry>,
       'dn': <distinguished name of entry>,
       'name': <ldap attribute associated with group_name_attr>
       '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 datetime import datetime
 
 
 from django.contrib import messages
 from django.contrib import messages
-from django.contrib.auth.models import User
 from django.contrib.sessions.models import Session
 from django.contrib.sessions.models import Session
 from django.db import DatabaseError
 from django.db import DatabaseError
 from django.db.models import Q
 from django.db.models import Q
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
 from desktop.auth.views import dt_logout
 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.models import UserProfile, get_profile
 from useradmin.views import import_ldap_users
 from useradmin.views import import_ldap_users
 
 
 from useradmin import ldap_access
 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__)
 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
 from nose.tools import assert_true, assert_equal, assert_false
 
 
 import desktop.conf
 import desktop.conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from django.conf import settings
 from django.conf import settings
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 from django.urls import reverse
 
 
 from useradmin.models import LdapGroup, UserProfile
 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 import ldap_access
 from useradmin.tests import BaseUserAdminTests, LdapTestConnection, reset_all_groups, reset_all_users
 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():
 def get_multi_ldap_config():
   return {'multi_ldap_conf': {
   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
 from nose.tools import assert_true, assert_equal, assert_false
 
 
 import desktop.conf
 import desktop.conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from django.conf import settings
 from django.conf import settings
-from django.contrib.auth.models import User, Group
 from django.urls import reverse
 from django.urls import reverse
 
 
 from useradmin.models import LdapGroup, UserProfile, get_profile
 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 import ldap_access
 from useradmin.tests import BaseUserAdminTests, LdapTestConnection, reset_all_groups, reset_all_users
 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):
 class TestUserAdminLdapDeprecated(BaseUserAdminTests):
   def test_useradmin_ldap_user_group_membership_sync(self):
   def test_useradmin_ldap_user_group_membership_sync(self):
     settings.MIDDLEWARE_CLASSES.append('useradmin.middleware.LdapSynchronizationMiddleware')
     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.plugins.skip import SkipTest
 from nose.tools import assert_true, assert_equal, assert_false, assert_not_equal
 from nose.tools import assert_true, assert_equal, assert_false, assert_not_equal
 from datetime import datetime
 from datetime import datetime
-from django.contrib.auth.models import User, Group
 from django.contrib.sessions.models import Session
 from django.contrib.sessions.models import Session
 from django.db.models import Q
 from django.db.models import Q
 from django.utils.encoding import smart_unicode
 from django.utils.encoding import smart_unicode
@@ -41,7 +40,7 @@ import desktop.conf
 
 
 from desktop import appmanager
 from desktop import appmanager
 from desktop.auth.backend import is_admin
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.views import home
 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.models import get_profile, get_default_user_group
 from useradmin.hue_password_policy import reset_password_policy
 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():
 def reset_all_users():
   """Reset to a clean state by deleting 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
 import json
 
 
 from nose.tools import assert_equal, assert_false, assert_true
 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
 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):
 class TestUseradminApi(object):
 
 

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

@@ -24,15 +24,12 @@ import threading
 import subprocess
 import subprocess
 import json
 import json
 
 
+
 from axes.decorators import FAILURE_LIMIT, LOCK_OUT_AT_FAILURE
 from axes.decorators import FAILURE_LIMIT, LOCK_OUT_AT_FAILURE
 from axes.models import AccessAttempt
 from axes.models import AccessAttempt
 from axes.utils import reset
 from axes.utils import reset
-
 import ldap
 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.urls import reverse
 from django.forms import ValidationError
 from django.forms import ValidationError
 from django.forms.utils import ErrorList
 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 _
 from django.utils.translation import get_language, ugettext as _
 
 
 import desktop.conf
 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.django_util import JsonResponse, render
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.models import _get_apps
 from desktop.models import _get_apps
 from desktop.views import antixss
 from desktop.views import antixss
-
 from hadoop.fs.exceptions import WebHdfsException
 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 HuePermission, UserProfile, LdapGroup
 from useradmin.models import get_profile, get_default_user_group
 from useradmin.models import get_profile, get_default_user_group
 from useradmin.forms import SyncLdapUsersGroupsForm, AddLdapGroupsForm, AddLdapUsersForm,\
 from useradmin.forms import SyncLdapUsersGroupsForm, AddLdapGroupsForm, AddLdapUsersForm,\
   PermissionsEditForm, GroupEditForm, SuperUserChangeForm, UserChangeForm, validate_username, validate_first_name, \
   PermissionsEditForm, GroupEditForm, SuperUserChangeForm, UserChangeForm, validate_username, validate_first_name, \
   validate_last_name, PasswordChangeForm
   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__)
 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 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.models import Document2
 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):
 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.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 from nose.plugins.skip import SkipTest
 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.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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.models import DocumentTag , Document
 from desktop.models import DocumentTag , Document
@@ -33,6 +31,11 @@ from desktop.models import DocumentTag , Document
 from pig.models import PigScript
 from pig.models import PigScript
 from useradmin.models import get_default_user_group
 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):
 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 import urlresolvers
 from django.core.exceptions import SuspiciousOperation
 from django.core.exceptions import SuspiciousOperation
 from django.contrib.auth import login, get_backends, authenticate
 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.contrib.sessions.models import Session
 from django.http import HttpResponseRedirect
 from django.http import HttpResponseRedirect
 from django.utils.translation import ugettext as _
 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.log.access import access_log, access_warn, last_access_map
 from desktop.settings import LOAD_BALANCER_COOKIE
 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:
 if sys.version_info[0] > 2:
   import urllib.request, urllib.error
   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_auth_ldap import backend as django_auth_ldap_backend
 from django.conf import settings
 from django.conf import settings
-from django.contrib.auth.models import User, Group
 from django.test.client import Client
 from django.test.client import Client
 
 
 from hadoop.test_base import PseudoHdfsTestBase
 from hadoop.test_base import PseudoHdfsTestBase
@@ -33,9 +32,15 @@ from useradmin.views import import_ldap_groups
 
 
 from desktop import conf, middleware
 from desktop import conf, middleware
 from desktop.auth import backend
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group
 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():
 def get_mocked_config():
   return {'mocked_ldap': {
   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 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.models import DefaultConfiguration
 from desktop.models import DefaultConfiguration
+
 from useradmin.models import get_default_user_group
 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):
 class TestDefaultConfiguration(object):
 
 

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

@@ -21,21 +21,26 @@ from datetime import datetime
 import json
 import json
 
 
 from nose.tools import assert_equal, assert_false, assert_true
 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.models import SavedQuery
 from beeswax.design import hql_query
 from beeswax.design import hql_query
+from librdbms.design import SQLdesign
 from oozie.models import Link, Workflow
 from oozie.models import Link, Workflow
 from oozie.tests import add_node
 from oozie.tests import add_node
 from pig.models import create_or_update_script
 from pig.models import create_or_update_script
 from useradmin.models import get_default_user_group
 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):
 class TestDocumentConverter(object):
 
 

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

@@ -23,10 +23,15 @@ import re
 import json
 import json
 
 
 import django.test.client
 import django.test.client
-from django.contrib.auth.models import User, Group
-
 import nose.tools
 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):
 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 logging
 import sys
 import sys
 
 
-from django.contrib.auth.models import User
+from urlparse import urlparse
 
 
 from desktop.auth.backend import is_admin
 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:
 if sys.version_info[0] > 2:
   from urllib.parse import urlparse as lib_urlparse
   from urllib.parse import urlparse as lib_urlparse
 else:
 else:
   from urlparse import urlparse as lib_urlparse
   from urlparse import urlparse as lib_urlparse
 
 
-LOG = logging.getLogger(__name__)
 
 
+LOG = logging.getLogger(__name__)
 DEFAULT_USER = DEFAULT_USER.get()
 DEFAULT_USER = DEFAULT_USER.get()
 
 
+
 class ProxyFS(object):
 class ProxyFS(object):
 
 
   def __init__(self, filesystems_dict, default_scheme, name='default'):
   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.tools import assert_raises, assert_false, eq_
 from nose import SkipTest
 from nose import SkipTest
 
 
-from django.contrib.auth.models import User
 from desktop.auth.backend import rewrite_user
 from desktop.auth.backend import rewrite_user
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.fs import ProxyFS
 from desktop.lib.fs import ProxyFS
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_permission, remove_from_group
 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():
 def test_fs_selection():
   make_logged_in_client(username='test', groupname='default', recreate=True, is_superuser=False)
   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 sys
 import time
 import time
 
 
-from django.contrib.auth.models import User
 from django.core.management.base import BaseCommand
 from django.core.management.base import BaseCommand
 from django.db import transaction
 from django.db import transaction
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.converters import DocumentConverter
 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):
 class Command(BaseCommand):
 
 
   def handle(self, *args, **options):
   def handle(self, *args, **options):

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

@@ -18,12 +18,17 @@
 
 
 import logging
 import logging
 
 
-from django.contrib.auth.models import User
 from django.core.management.base import CommandError, BaseCommand
 from django.core.management.base import CommandError, BaseCommand
 from django.utils.translation import ugettext_lazy as _
 from django.utils.translation import ugettext_lazy as _
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Document2
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -24,10 +24,14 @@ import logging
 import multiprocessing
 import multiprocessing
 import threading
 import threading
 
 
-from django.contrib.auth.models import User
-
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.metrics import global_registry
 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__)
 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 import messages
 from django.contrib.auth import REDIRECT_FIELD_NAME, BACKEND_SESSION_KEY, authenticate, load_backend, login
 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.middleware import RemoteUserMiddleware
-from django.contrib.auth.models import User
 from django.core import exceptions, urlresolvers
 from django.core import exceptions, urlresolvers
 from django.http import HttpResponseNotAllowed, HttpResponseForbidden
 from django.http import HttpResponseNotAllowed, HttpResponseForbidden
 from django.urls import resolve
 from django.urls import resolve
@@ -50,6 +49,7 @@ from hadoop import cluster
 import desktop.views
 import desktop.views
 import desktop.conf
 import desktop.conf
 from desktop.auth.backend import is_admin
 from desktop.auth.backend import is_admin
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.context_processors import get_app_name
 from desktop.context_processors import get_app_name
 from desktop.lib import apputil, i18n, fsmanager
 from desktop.lib import apputil, i18n, fsmanager
 from desktop.lib.django_util import JsonResponse, render, render_json
 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.log.access import access_log, log_page_hit, access_warn
 from desktop import appmanager
 from desktop import appmanager
 from desktop import metrics
 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__)
 LOG = logging.getLogger(__name__)

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

@@ -17,11 +17,16 @@
 
 
 import re
 import re
 
 
-from django.contrib.auth.models import User
 from django.core.validators import RegexValidator
 from django.core.validators import RegexValidator
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import get_username_re_rule
 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():
 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.management import call_command
 from django.core.paginator import Paginator
 from django.core.paginator import Paginator
 from django.conf.urls import url
 from django.conf.urls import url
-from django.contrib.auth.models import User
 from django.db import connection
 from django.db import connection
 from django.urls import reverse
 from django.urls import reverse
 from django.test.client import Client
 from django.test.client import Client
@@ -60,7 +59,7 @@ import desktop.views as views
 
 
 from desktop.auth.backend import rewrite_user
 from desktop.auth.backend import rewrite_user
 from desktop.appmanager import DESKTOP_APPS
 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.django_test_util import make_logged_in_client
 from desktop.lib.conf import validate_path
 from desktop.lib.conf import validate_path
 from desktop.lib.django_util import TruncatingModel
 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.paths import get_desktop_root
 from desktop.lib.python_util import force_dict_to_strings
 from desktop.lib.python_util import force_dict_to_strings
 from desktop.lib.test_utils import grant_access
 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.models import Directory, Document, Document2, get_data_link, _version_from_properties, ClusterConfig, HUE_VERSION
 from desktop.redaction import logfilter
 from desktop.redaction import logfilter
 from desktop.redaction.engine import RedactionPolicy, RedactionRule
 from desktop.redaction.engine import RedactionPolicy, RedactionRule
@@ -80,8 +80,15 @@ if sys.version_info[0] > 2:
 else:
 else:
   from cStringIO import StringIO as string_io
   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__)
 LOG = logging.getLogger(__name__)
 
 
+
 def test_home():
 def test_home():
   c = make_logged_in_client(username="test_home", groupname="test_home", recreate=True, is_superuser=False)
   c = make_logged_in_client(username="test_home", groupname="test_home", recreate=True, is_superuser=False)
   user = User.objects.get(username="test_home")
   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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_false, assert_true, assert_not_equal, assert_raises
 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.core import management
 from django.db.utils import OperationalError
 from django.db.utils import OperationalError
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.converters import DocumentConverter
 from desktop.converters import DocumentConverter
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.fs import ProxyFS
 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 useradmin.models import get_default_user_group
 from oozie.models2 import Workflow
 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):
 class MockFs(object):
   def __init__(self):
   def __init__(self):

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

@@ -20,9 +20,9 @@ import os
 import tempfile
 import tempfile
 import string
 import string
 
 
-from django.contrib.auth.models import User
 from nose.tools import assert_equal, assert_false, assert_true, assert_raises, eq_
 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.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
 from desktop.lib.test_utils import grant_access, add_to_group, add_permission, remove_from_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.s3test_utils import S3TestBase, generate_id
 from aws.s3.upload import DEFAULT_WRITE_SIZE
 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):
 class S3FSTest(S3TestBase):
 
 
@@ -288,7 +293,7 @@ class S3FSTest(S3TestBase):
   def test_mkdir(self):
   def test_mkdir(self):
     dir_path = self.get_test_path('test_mkdir')
     dir_path = self.get_test_path('test_mkdir')
     assert_false(self.fs.exists(dir_path))
     assert_false(self.fs.exists(dir_path))
-    
+
     self.fs.mkdir(dir_path)
     self.fs.mkdir(dir_path)
     assert_true(self.fs.exists(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
 from builtins import object
 import json
 import json
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 from desktop.lib.rest import resource
 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.models import Collection2, augment_response
 from dashboard.controller import DashboardController
 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}
 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 logging
 import os
 import os
 
 
-from django.contrib.auth.models import User
 from django.core.urlresolvers import reverse
 from django.core.urlresolvers import reverse
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from hadoop.fs.hadoopfs import Hdfs
 from hadoop.fs.hadoopfs import Hdfs
 from notebook.models import make_notebook
 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 indexer.conf import CONFIG_JARS_LIBS_PATH, config_morphline_path
 from libzookeeper.conf import zkensemble
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

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

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

@@ -18,16 +18,21 @@ from builtins import object
 import logging
 import logging
 import os
 import os
 
 
-from django.contrib.auth.models import User
 from django.core.urlresolvers import reverse
 from django.core.urlresolvers import reverse
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from libzookeeper.conf import zkensemble
 from libzookeeper.conf import zkensemble
 from metadata.manager_client import ManagerApi
 from metadata.manager_client import ManagerApi
 
 
 from indexer.conf import config_morphline_path
 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__)
 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
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # limitations under the License.
 
 
-from django.contrib.auth.models import User
-
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 from nose.tools import assert_equal, assert_true
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
+
 from indexer.indexers.flume import FlumeIndexer
 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():
 def test_generate_from_directory_to_solr_index():
   raise SkipTest
   raise SkipTest
@@ -154,7 +159,7 @@ morphlines : [
     ,
     ,
     configs[0][1].strip() # 'agent_morphlines_conf_file'
     configs[0][1].strip() # 'agent_morphlines_conf_file'
   )
   )
-  
+
   assert_equal(
   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')
     ('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 collections import deque
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 from mako.lookup import TemplateLookup
 from mako.lookup import TemplateLookup
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Document2
 from desktop.models import Document2
 from notebook.connectors.base import get_api
 from notebook.connectors.base import get_api
 from notebook.models import Notebook, make_notebook
 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.fields import get_field_type
 from indexer.file_format import get_file_format_instance, get_file_format_class
 from indexer.file_format import get_file_format_instance, get_file_format_class
 from indexer.indexers.morphline_operations import get_checked_args
 from indexer.indexers.morphline_operations import get_checked_args
 from indexer.solr_client import SolrClient
 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__)
 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.attrib import attr
 from nose.plugins.skip import SkipTest
 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.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
 from hadoop.pseudo_hdfs4 import is_live_cluster, shared_cluster
 from hadoop.pseudo_hdfs4 import is_live_cluster, shared_cluster
@@ -48,6 +47,12 @@ if sys.version_info[0] > 2:
 else:
 else:
   from StringIO import StringIO as string_io
   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__)
 LOG = logging.getLogger(__name__)
 
 
 
 

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

@@ -19,10 +19,10 @@ import json
 import logging
 import logging
 import uuid
 import uuid
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.i18n import smart_str
 from desktop.lib.i18n import smart_str
 from librdbms.conf import DATABASES, get_database_password, get_server_choices, get_connector_name
 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.connectors.base import get_api
 from notebook.models import make_notebook
 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__)
 LOG = logging.getLogger(__name__)
 
 
@@ -271,4 +276,4 @@ def _splitby_column_check(statement, destination_splitby_column):
       'statement': statement,
       'statement': statement,
       'destinationSplitbyColumn': destination_splitby_column
       '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 urllib.request, urllib.error
 import uuid
 import uuid
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.translation import ugettext as _
 from django.utils.translation import ugettext as _
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib import django_mako
 from desktop.lib import django_mako
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 
 
@@ -38,6 +38,12 @@ else:
   from urllib import unquote as urllib_unquote
   from urllib import unquote as urllib_unquote
   from urlparse import urlparse
   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__)
 LOG = logging.getLogger(__name__)
 
 
 
 

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

@@ -19,14 +19,19 @@
 from builtins import object
 from builtins import object
 import json
 import json
 
 
-from django.contrib.auth.models import User
 from mock import patch, Mock, MagicMock
 from mock import patch, Mock, MagicMock
 from nose.tools import assert_equal, assert_true
 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 desktop.lib.django_test_util import make_logged_in_client
 
 
 from indexer.indexers.sql import SQLIndexer
 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):
 class MockRequest(object):
   def __init__(self, fs=None, user=None):
   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
 from builtins import object
 import logging
 import logging
 
 
-from django.contrib.auth.models import User
-
 from nose.plugins.skip import SkipTest
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_false, assert_not_equal, assert_true
 from nose.tools import assert_equal, assert_false, assert_not_equal, assert_true
 
 
 from desktop.auth.backend import rewrite_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.django_test_util import make_logged_in_client
 
 
 from indexer.conf import ENABLE_SQOOP
 from indexer.conf import ENABLE_SQOOP
 from indexer.indexers.rdbms import _get_api
 from indexer.indexers.rdbms import _get_api
 from librdbms.server import dbms as rdbms
 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__)
 LOG = logging.getLogger(__name__)
 
 
+
 class TestRdbmsIndexer(object):
 class TestRdbmsIndexer(object):
 
 
   @classmethod
   @classmethod
@@ -62,4 +67,4 @@ class TestRdbmsIndexer(object):
     assert_not_equal('', data['rows'], data)
     assert_not_equal('', data['rows'], data)
 
 
 class Bag(dict):
 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 builtins import object
 from nose.tools import assert_equal, assert_true, assert_false
 from nose.tools import assert_equal, assert_true, assert_false
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 
 
 from indexer.solr_client import SolrClient
 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):
 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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false
 from nose.tools import assert_equal, assert_true, assert_false
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster, get_db_prefix
 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.conf import get_solr_ensemble
 from indexer.controller import CollectionManagerController
 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():
 def test_get_ensemble():
 
 

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

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

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

@@ -21,7 +21,7 @@ try:
   import oauth2 as oauth
   import oauth2 as oauth
 except:
 except:
   oauth = None
   oauth = None
- 
+
 import logging
 import logging
 import urllib.request, urllib.parse, urllib.error
 import urllib.request, urllib.parse, urllib.error
 import httplib2
 import httplib2
@@ -30,7 +30,6 @@ import django.contrib.auth.views
 from django.core import urlresolvers
 from django.core import urlresolvers
 from django.core.exceptions import SuspiciousOperation
 from django.core.exceptions import SuspiciousOperation
 from django.contrib.auth import login, get_backends, authenticate
 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.contrib.sessions.models import Session
 from django.http import HttpResponseRedirect
 from django.http import HttpResponseRedirect
 from django.utils.translation import ugettext as _
 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.backend import AllowFirstUserDjangoBackend
 from desktop.auth.forms import UserCreationForm, AuthenticationForm
 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 render
 from desktop.lib.django_util import login_notrequired
 from desktop.lib.django_util import login_notrequired
 from desktop.log.access import access_warn, last_access_map
 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
 import liboauth.conf
 from liboauth.backend import OAuthBackend
 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
 @login_notrequired
 def show_login_page(request, login_errors=False):
 def show_login_page(request, login_errors=False):
@@ -77,10 +82,10 @@ def oauth_login(request):
 
 
   return HttpResponseRedirect(url)
   return HttpResponseRedirect(url)
 
 
-  
+
 @login_notrequired
 @login_notrequired
 def oauth_authenticated(request):
 def oauth_authenticated(request):
-   
+
   access_token, next = OAuthBackend.handleAuthenticationRequest(request)
   access_token, next = OAuthBackend.handleAuthenticationRequest(request)
   if access_token == "":
   if access_token == "":
       return show_login_page(request, True)
       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
 from builtins import object
 import logging
 import logging
 
 
-from django.contrib.auth.models import User
 from nose.plugins.attrib import attr
 from nose.plugins.attrib import attr
 from nose.tools import assert_equal, assert_true, assert_not_equal, assert_raises
 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 import cluster, pseudo_hdfs4
 from hadoop.conf import HDFS_CLUSTERS, MR_CLUSTERS, YARN_CLUSTERS
 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.test_utils import clear_sys_caches
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
@@ -40,6 +40,11 @@ from liboozie.credentials import Credentials
 from liboozie.credentials_tests import TestCredentials
 from liboozie.credentials_tests import TestCredentials
 from liboozie.submission2 import Submission
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -18,13 +18,13 @@
 from builtins import object
 from builtins import object
 import logging
 import logging
 
 
-from django.contrib.auth.models import User
 from nose.plugins.attrib import attr
 from nose.plugins.attrib import attr
 from nose.tools import assert_equal, assert_true, assert_not_equal
 from nose.tools import assert_equal, assert_true, assert_not_equal
 
 
 from hadoop import cluster, pseudo_hdfs4
 from hadoop import cluster, pseudo_hdfs4
 from hadoop.conf import HDFS_CLUSTERS, MR_CLUSTERS, YARN_CLUSTERS
 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.test_utils import clear_sys_caches
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from oozie.tests import OozieMockBase
 from oozie.tests import OozieMockBase
@@ -32,6 +32,11 @@ from useradmin.views import ensure_home_directory
 
 
 from liboozie.submittion import Submission
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -17,10 +17,10 @@
 
 
 import logging
 import logging
 
 
-from django.contrib.auth.models import User
 from nose.tools import assert_equal
 from nose.tools import assert_equal
 
 
 import desktop.conf as desktop_conf
 import desktop.conf as desktop_conf
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.lib.test_utils import reformat_xml
 from desktop.lib.test_utils import reformat_xml
 
 
 from liboozie import conf
 from liboozie import conf
@@ -28,6 +28,11 @@ from liboozie.types import WorkflowAction, Coordinator
 from liboozie.utils import config_gen
 from liboozie.utils import config_gen
 from oozie.tests import MockOozieApi
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -23,14 +23,19 @@ import logging
 import sys
 import sys
 
 
 from django.contrib.auth import logout as auth_logout
 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 django_openid_auth.auth import OpenIDBackend as _OpenIDBackend
 
 
 from desktop.auth.backend import rewrite_user
 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 useradmin.models import get_profile, get_default_user_group, UserProfile
 
 
 from libopenid import metrics
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -23,18 +23,23 @@ from __future__ import absolute_import
 import logging
 import logging
 
 
 from django.contrib.auth import logout as auth_logout
 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.backends import Saml2Backend as _Saml2Backend
 from djangosaml2.views import logout as saml_logout
 from djangosaml2.views import logout as saml_logout
 
 
 from desktop.auth.backend import force_username_case, rewrite_user
 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 conf
 from libsaml import metrics
 from libsaml import metrics
 
 
 from useradmin.models import get_profile, get_default_user_group, UserProfile
 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__)
 LOG = logging.getLogger(__name__)
 
 
 
 

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

@@ -19,14 +19,19 @@
 from builtins import object
 from builtins import object
 import pickle
 import pickle
 
 
-from django.contrib.auth.models import User
 from nose.tools import assert_equal, assert_false, assert_true
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
 from desktop.lib.test_utils import grant_access
 
 
 from libsentry.privilege_checker import PrivilegeChecker
 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):
 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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false, assert_not_equal, assert_raises
 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.django_test_util import make_logged_in_client
 from desktop.lib.exceptions import StructuredThriftTransportException
 from desktop.lib.exceptions import StructuredThriftTransportException
 from desktop.lib.exceptions_renderable import PopupException
 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_ha import get_next_available_server
 from libsentry.sentry_site import get_sentry_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):
 def create_mock_client_fn(client_class, username, server, component=None):
   class MockSentryClient(object):
   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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 from nose.tools import assert_equal, assert_true
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 from dashboard.models import Collection2
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
 from hadoop.pseudo_hdfs4 import is_live_cluster
 
 
 from libsolr.api import SolrApi
 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__)
 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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
 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.models import ZookeeperClient
 from libzookeeper.conf import zkensemble, ENSEMBLE
 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):
 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 nose.tools import assert_equal
 
 
 from django.core.cache import cache
 from django.core.cache import cache
-from django.contrib.auth.models import User
 
 
 from desktop.auth.backend import rewrite_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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 from libsentry.privilege_checker import SENTRY_PRIVILEGE_CACHE_KEY, PrivilegeChecker
 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.metadata_sites import get_navigator_hue_server_name
 from metadata.catalog.navigator_client import NavigatorApi
 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__)
 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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 from nose.tools import assert_equal, assert_true
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 
 
 from desktop.auth.backend import rewrite_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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
 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.navigator_client import NavigatorApi
 from metadata.catalog.atlas_client import AtlasApi
 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__)
 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 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.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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 
 
 from metadata.optimizer_api import _convert_queries
 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__)
 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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true, assert_false
 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.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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 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 metadata.conf import OPTIMIZER, has_optimizer
 from desktop.models import uuid_default
 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__)
 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.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_true
 from nose.tools import assert_equal, assert_true
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from TCLIService.ttypes import TStatusCode, TProtocolVersion, TOperationType
 from TCLIService.ttypes import TStatusCode, TProtocolVersion, TOperationType
 
 
 from desktop.auth.backend import rewrite_user
 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.i18n import smart_str
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 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.server import dbms
 from beeswax.test_base import BeeswaxSampleProvider, get_query_server_config, is_hive_on_spark
 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__)
 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 mock import patch, Mock, MagicMock
 from nose.tools import assert_equal, assert_not_equal, assert_true, assert_false
 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.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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import add_to_group, grant_access
 from desktop.lib.test_utils import add_to_group, grant_access
 
 
 from notebook.connectors.sql_alchemy import SqlAlchemyApi
 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__)
 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 os
 import pwd
 import pwd
 
 
-from django.contrib.auth.models import User
 from django.core import management
 from django.core import management
 from django.core.management.base import BaseCommand
 from django.core.management.base import BaseCommand
 from django.db import transaction
 from django.db import transaction
 
 
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.models import Directory, Document, Document2, Document2Permission, SAMPLE_USER_OWNERS
 from desktop.models import Directory, Document, Document2, Document2Permission, SAMPLE_USER_OWNERS
 from useradmin.models import get_default_user_group, install_sample_user
 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__)
 LOG = logging.getLogger(__name__)
 
 

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

@@ -29,14 +29,13 @@ import uuid
 
 
 from datetime import timedelta
 from datetime import timedelta
 
 
-from django.contrib.auth.models import User
 from django.contrib.sessions.models import Session
 from django.contrib.sessions.models import Session
 from django.db.models import Count
 from django.db.models import Count
 from django.db.models.functions import Trunc
 from django.db.models.functions import Trunc
 from django.utils.html import escape
 from django.utils.html import escape
 from django.utils.translation import ugettext as _
 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.i18n import smart_unicode
 from desktop.lib.paths import SAFE_CHARACTERS_URI
 from desktop.lib.paths import SAFE_CHARACTERS_URI
 from desktop.models import Document2
 from desktop.models import Document2
@@ -49,6 +48,12 @@ if sys.version_info[0] > 2:
 else:
 else:
   from urllib import quote as urllib_quote
   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__)
 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.cache import caches
 from django.core.files.storage import get_storage_class
 from django.core.files.storage import get_storage_class
-from django.contrib.auth.models import User
 from django.db import transaction
 from django.db import transaction
 from django.http import FileResponse, HttpRequest
 from django.http import FileResponse, HttpRequest
 
 
 from beeswax import data_export
 from beeswax import data_export
 from desktop.auth.backend import rewrite_user
 from desktop.auth.backend import rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 from desktop.celery import app
 from desktop.celery import app
 from desktop.conf import TASK_SERVER
 from desktop.conf import TASK_SERVER
 from desktop.lib import export_csvxls
 from desktop.lib import export_csvxls
@@ -52,6 +52,12 @@ if sys.version_info[0] > 2:
 else:
 else:
   from StringIO import StringIO as string_io
   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_TASK = get_task_logger(__name__)
 LOG = logging.getLogger(__name__)
 LOG = logging.getLogger(__name__)
 STATE_MAP = {
 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.plugins.attrib import attr
 from nose.tools import assert_equal, assert_true, assert_false
 from nose.tools import assert_equal, assert_true, assert_false
 
 
-from django.contrib.auth.models import User
 from django.urls import reverse
 from django.urls import reverse
 from azure.conf import is_adls_enabled
 from azure.conf import is_adls_enabled
 
 
 from desktop import appmanager
 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.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access, add_permission
 from desktop.lib.test_utils import grant_access, add_permission
 from desktop.models import Directory, Document, Document2
 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.conf import get_ordered_interpreters, INTERPRETERS_SHOWN_ON_WHEEL, INTERPRETERS
 from notebook.models import Analytics
 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):
 class TestNotebookApi(object):