jumpserver/apps/settings/utils/db.py

55 lines
1.6 KiB
Python

from django.apps import apps
from orgs.utils import tmp_to_root_org
__all__ = ['DBTableDataAnalyzer']
class DBTableDataAnalyzer(object):
def get_tables_info(self):
table_models = self.get_all_table_models()
with tmp_to_root_org():
table_info = self._fetch_tables_info(table_models)
return table_info
def get_all_table_models(self):
# construct tables models
table_models = set()
for model in apps.get_models():
table_models.add(model)
for m2m_field in model._meta.many_to_many:
table_models.add(m2m_field.remote_field.through)
return table_models
def _fetch_tables_info(self, table_models):
"""
return:
info = {
'table_name': {
'count': 100
}
...
}
"""
info = {}
for table_model in table_models:
if table_model._meta.proxy == True:
# skip proxy model
continue
count = self._fetch_row_count(table_model)
data = {
"count": count
}
table_name = table_model._meta.db_table
info.update({f"{table_name}": data})
return info
def _fetch_row_count(self, table_model):
# fetch table row count
if hasattr(table_model, 'objects_raw'):
# use objects_raw to count
count = table_model.objects_raw.count()
else:
count = table_model.objects.count()
return count