codekingpro/portable-devtools

1##########################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8##########################################################################9 10from pgadmin.browser.server_groups.servers.databases.schemas.tables.utils \11    import BaseTableView12from pgadmin.browser.server_groups.servers.databases.schemas.utils \13    import get_schemas14from pgadmin.browser.server_groups.servers.databases.schemas.utils \15    import DataTypeReader16from pgadmin.utils.preferences import Preferences17 18 19class ERDTableView(BaseTableView, DataTypeReader):20    def __init__(self):21        super().__init__(cmd='erd')22 23    @BaseTableView.check_precondition24    def sql(self, conn_id=None, did=None, sid=None, data={}, with_drop=False):25        return BaseTableView.get_sql(self, did, None, None, data, None,26                                     add_not_exists_clause=True,27                                     with_drop=with_drop)28 29    @BaseTableView.check_precondition30    def get_types(self, conn_id=None, did=None, sid=None):31        condition = self.get_types_condition_sql(False)32        return DataTypeReader.get_types(self, self.conn, condition, True)33 34    @BaseTableView.check_precondition35    def fetch_all_tables(self, did=None, sid=None, scid=None):36        all_tables = []37        schemas = {'rows': []}38        if scid is None:39            status, schemas = get_schemas(self.conn, show_system_objects=False)40            if not status:41                return status, schemas42        else:43            schemas['rows'].append({'oid': scid})44 45        for row in schemas['rows']:46            status, res = \47                BaseTableView.fetch_tables(self, sid, did, row['oid'],48                                           with_serial_cols=True)49            if not status:50                return status, res51 52            all_tables.extend(res.values())53 54        return True, all_tables55 56    @BaseTableView.check_precondition57    def traverse_related_tables(self, did=None, sid=None, scid=None,58                                tid=None, related={}, maxdepth=0, currdepth=0):59 60        status, res = \61            BaseTableView.fetch_tables(self, sid, did, scid, tid=tid,62                                       with_serial_cols=True)63 64        if not status:65            return status, res66 67        related[tid] = res68        # Max depth limit reached69        if currdepth == maxdepth:70            new_fks = []71            for fk in related[tid].pop('foreign_key', []):72                if fk['confrelid'] in related:73                    new_fks.append(fk)74 75            related[tid]['foreign_key'] = new_fks76            return True, None77 78        status, depending_res = BaseTableView.get_fk_ref_tables(79            self, tid)80 81        if not status:82            return status, depending_res83 84        for fk in [*res.get('foreign_key', []), *depending_res]:85            if fk['confrelid'] in related:86                continue87            status, res = self.traverse_related_tables(88                did=did, sid=sid, scid=fk['refnspoid'], tid=fk['confrelid'],89                related=related, maxdepth=maxdepth, currdepth=currdepth + 1)90            if not status:91                return status, res92 93        return True, None94 95 96class ERDHelper:97    def __init__(self, conn_id, sid, did):98        self.conn_id = conn_id99        self.did = did100        self.sid = sid101        self.table_view = ERDTableView()102        self.link_view = None103 104    def get_types(self):105        return self.table_view.get_types(106            conn_id=self.conn_id, did=self.did, sid=self.sid)107 108    def get_table_sql(self, data, with_drop=False):109        SQL, _ = self.table_view.sql(110            conn_id=self.conn_id, did=self.did, sid=self.sid,111            data=data, with_drop=with_drop)112        return SQL113 114    def get_all_tables(self, scid, tid):115        if tid is None and scid is None:116            status, res = self.table_view.fetch_all_tables(117                did=self.did, sid=self.sid)118        elif tid is None:119            status, res = self.table_view.fetch_all_tables(120                did=self.did, sid=self.sid, scid=scid)121        else:122            prefs = Preferences.module('erd')123            table_relation_depth = prefs.preference('table_relation_depth')124            related = {}125            status, res = self.table_view.traverse_related_tables(126                did=self.did, sid=self.sid, scid=scid, tid=tid,127                related=related, maxdepth=table_relation_depth.get()128            )129            if status:130                res = list(related.values())131        return status, res132