feat(ingest/postgres): fetch table sizes for profile (datahub-project…

…#10864)
acryldata · Jul 16, 2024 · 298c299 · 298c299
1 parent 4b83adf
commit 298c299
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/postgres.py b/metadata-ingestion/src/datahub/ingestion/source/sql/postgres.py
@@ -276,3 +276,19 @@ def get_identifier(
             return f"{self.config.database}.{regular}"
         current_database = self.get_db_name(inspector)
         return f"{current_database}.{regular}"
+
+    def add_profile_metadata(self, inspector: Inspector) -> None:
+        try:
+            with inspector.engine.connect() as conn:
+                for row in conn.execute(
+                    """SELECT table_catalog, table_schema, table_name, pg_table_size('"' || table_catalog || '"."' || table_schema || '"."' || table_name || '"') AS table_size FROM information_schema.TABLES"""
+                ):
+                    self.profile_metadata_info.dataset_name_to_storage_bytes[
+                        self.get_identifier(
+                            schema=row.table_schema,
+                            entity=row.table_name,
+                            inspector=inspector,
+                        )
+                    ] = row.table_size
+        except Exception as e:
+            logger.error(f"failed to fetch profile metadata: {e}")
diff --git a/metadata-ingestion/tests/integration/postgres/postgres_all_db_mces_with_db_golden.json b/metadata-ingestion/tests/integration/postgres/postgres_all_db_mces_with_db_golden.json
@@ -832,7 +832,8 @@
                 {
                     "fieldPath": "metadata_json"
                 }
-            ]
+            ],
+            "sizeInBytes": 16384
         }
     },
     "systemMetadata": {

diff --git a/metadata-ingestion/tests/integration/postgres/postgres_mces_with_db_golden.json b/metadata-ingestion/tests/integration/postgres/postgres_mces_with_db_golden.json
@@ -600,7 +600,8 @@
             },
             "rowCount": 2,
             "columnCount": 9,
-            "fieldProfiles": []
+            "fieldProfiles": [],
+            "sizeInBytes": 16384
         }
     },
     "systemMetadata": {