From 5a57c296a1ddee397ddbabc555ed1f632b23430f Mon Sep 17 00:00:00 2001
From: Bai <baijiangjie@gmail.com>
Date: Tue, 9 Jul 2024 18:44:19 +0800
Subject: [PATCH] feat: add db table data analyzer util

---
 apps/authentication/models/__init__.py |  2 +
 apps/settings/utils/__init__.py        |  1 +
 apps/settings/utils/db.py              | 55 ++++++++++++++++++++++++++
 3 files changed, 58 insertions(+)
 create mode 100644 apps/settings/utils/db.py

diff --git a/apps/authentication/models/__init__.py b/apps/authentication/models/__init__.py
index e889e03b1..23f8838ec 100644
--- a/apps/authentication/models/__init__.py
+++ b/apps/authentication/models/__init__.py
@@ -3,3 +3,5 @@ from .connection_token import *
 from .private_token import *
 from .sso_token import *
 from .temp_token import *
+from ..backends.passkey.models import *
+
diff --git a/apps/settings/utils/__init__.py b/apps/settings/utils/__init__.py
index 91d05860e..f2215e770 100644
--- a/apps/settings/utils/__init__.py
+++ b/apps/settings/utils/__init__.py
@@ -3,3 +3,4 @@
 
 from .ldap import *
 from .common import *
+from .db import *
diff --git a/apps/settings/utils/db.py b/apps/settings/utils/db.py
new file mode 100644
index 000000000..8c2fca321
--- /dev/null
+++ b/apps/settings/utils/db.py
@@ -0,0 +1,55 @@
+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
+    
\ No newline at end of file