feat(ingest/bigquery): improve profiler to support multiple partition columns and support external table profiling

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -62,20 +62,19 @@ class BigqueryTableConstraint:
 
 @dataclass
 class PartitionInfo:
-    field: str
-    # Data type is optional as we not have it when we set it from TimePartitioning
-    column: Optional[BigqueryColumn] = None
+    fields: List[str]  # Changed from single field to list
+    columns: Optional[List[BigqueryColumn]] = None
     type: str = TimePartitioningType.DAY
     expiration_ms: Optional[int] = None
     require_partition_filter: bool = False
 
-    # TimePartitioning field doesn't provide data_type so we have to add it afterwards
     @classmethod
     def from_time_partitioning(
         cls, time_partitioning: TimePartitioning
     ) -> "PartitionInfo":
+        """Convert BigQuery time partitioning to PartitionInfo."""
         return cls(
-            field=time_partitioning.field or "_PARTITIONTIME",
+            fields=[time_partitioning.field or "_PARTITIONTIME"],  # Now a list
             type=time_partitioning.type_,
             expiration_ms=time_partitioning.expiration_ms,
             require_partition_filter=time_partitioning.require_partition_filter,
@@ -90,7 +89,7 @@ def from_range_partitioning(
             return None
 
         return cls(
-            field=field,
+            fields=[field],
             type=RANGE_PARTITION_NAME,
         )
 

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -630,15 +630,13 @@ def _process_table(
             )
 
         # If table has time partitioning, set the data type of the partitioning field
-        if table.partition_info:
-            table.partition_info.column = next(
-                (
-                    column
-                    for column in columns
-                    if column.name == table.partition_info.field
-                ),
-                None,
-            )
+        if table.partition_info and table.partition_info.fields:
+            matching_columns = [
+                column
+                for column in columns
+                if column.name in table.partition_info.fields
+            ]
+            table.partition_info.columns = matching_columns
         yield from self.gen_table_dataset_workunits(
             table, columns, project_id, dataset_name
         )