doc_count_util.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  1. # Licensed to Cloudera, Inc. under one
  2. # or more contributor license agreements. See the NOTICE file
  3. # distributed with this work for additional information
  4. # regarding copyright ownership. Cloudera, Inc. licenses this file
  5. # to you under the Apache License, Version 2.0 (the
  6. # "License"); you may not use this file except in compliance
  7. # with the License. You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. import json
  17. import logging
  18. import time
  19. import re
  20. import sys
  21. from django.db import transaction
  22. from desktop.models import Document, DocumentPermission, DocumentTag, Document2, Directory, Document2Permission
  23. LOG = logging.getLogger(__name__)
  24. class DocumentCounts(object):
  25. """
  26. Given a user, prints the count of doc1 and doc2 documents
  27. """
  28. def __init__(self, user, differ=False):
  29. LOG.info("Sai")
  30. self.user = user
  31. self.differ = differ
  32. self.home_dir = Document2.objects.create_user_directories(self.user)
  33. sys.stdout = open('counts.txt', 'a')
  34. def printCounts(self):
  35. LOG.info("HI")
  36. # Convert SavedQuery documents
  37. try:
  38. from beeswax.models import SavedQuery, HQL, IMPALA, RDBMS
  39. docs = self._get_unconverted_docs(SavedQuery).filter(extra__in=[HQL, IMPALA, RDBMS])
  40. self.saved_query_doc1 = len(docs)
  41. docs = Document2.objects.filter(owner=self.user, type__startswith='query-').filter(is_history=False)
  42. self.saved_query_doc2 = len(docs)
  43. except:
  44. LOG.info('Cannot convert Saved Query documents: beeswax app is not installed')
  45. pass
  46. # Convert SQL Query history documents
  47. try:
  48. from beeswax.models import SavedQuery, HQL, IMPALA, RDBMS
  49. docs = self._get_unconverted_docs(SavedQuery, with_history=True).filter(extra__in=[HQL, IMPALA, RDBMS]).order_by('-last_modified')
  50. self.query_history_doc1 = len(docs)
  51. docs = Document2.objects.filter(owner=self.user,type__startswith='query-').filter(is_history=True)
  52. self.query_history_doc2 = len(docs)
  53. except:
  54. LOG.info('Cannot convert Saved Query documents: beeswax app is not installed')
  55. pass
  56. if self.differ:
  57. if self.saved_query_doc2 != self.saved_query_doc1 or self.query_history_doc2 != self.query_history_doc1:
  58. print("%-20s %s/%s \t %s/%s" %(self.user.username, self.saved_query_doc2, self.saved_query_doc1, self.query_history_doc2, self.query_history_doc1))
  59. else:
  60. print("%-20s %s/%s \t %s/%s" %(self.user.username, self.saved_query_doc2, self.saved_query_doc1, self.query_history_doc2, self.query_history_doc1))
  61. sys.stdout.close()
  62. def _get_unconverted_docs(self, content_type, with_history=False):
  63. docs = Document.objects.get_docs(self.user, content_type).filter(owner=self.user)
  64. tags = [
  65. DocumentTag.objects.get_trash_tag(user=self.user), # No trashed docs
  66. DocumentTag.objects.get_example_tag(user=self.user), # No examples
  67. # self.imported_tag # No already imported docs
  68. ]
  69. if not with_history:
  70. tags.append(DocumentTag.objects.get_history_tag(user=self.user)) # No history yet
  71. return docs.exclude(tags__in=tags)
  72. def _get_parent_directory(self, document):
  73. """
  74. Returns the parent directory object that should be used for a given document. If the document is tagged with a
  75. project name (non-RESERVED DocumentTag), a Directory object with the first project tag found is returned.
  76. Otherwise, the owner's home directory is returned.
  77. """
  78. parent_dir = self.home_dir
  79. project_tags = document.tags.exclude(tag__in=DocumentTag.RESERVED)
  80. if project_tags.exists():
  81. first_tag = project_tags[0]
  82. parent_dir, created = Directory.objects.get_or_create(
  83. owner=self.user,
  84. name=first_tag.tag,
  85. parent_directory=self.home_dir
  86. )
  87. return parent_dir