From 0ff71e69ac91d8ff9f695909099981555d72f4b0 Mon Sep 17 00:00:00 2001
From: Ravi Suhag <suhag.ravi@gmail.com>
Date: Wed, 29 May 2024 12:19:56 -0500
Subject: [PATCH] feat(bigquery): adds sql query for VIEW and MATERIALIZED VIEW
 in asset for BQ (#488)

* feat: group comamnds

* feat(bigquery): adds sql query for VIEW and MATERIALIZED VIEW in asset for BQ
---
 plugins/extractors/bigquery/bigquery.go | 57 +++++++++++++------------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/plugins/extractors/bigquery/bigquery.go b/plugins/extractors/bigquery/bigquery.go
index fc4448758..27047644e 100755
--- a/plugins/extractors/bigquery/bigquery.go
+++ b/plugins/extractors/bigquery/bigquery.go
@@ -415,44 +415,47 @@ func (e *Extractor) buildAsset(ctx context.Context, t *bigquery.Table, md *bigqu
 		}
 	}
 
-	table, err := anypb.New(&v1beta2.Table{
-		Columns:       e.buildColumns(ctx, md.Schema, md),
-		PreviewFields: previewFields,
-		PreviewRows:   previewRows,
-		Profile:       tableProfile,
-		Attributes: utils.TryParseMapToProto(map[string]interface{}{
-			"full_qualified_name": tableFQN,
-			"dataset":             t.DatasetID,
-			"project":             t.ProjectID,
-			"type":                string(md.Type),
-			"partition_data":      partitionData,
-			"clustering_fields":   clusteringFields,
-		}),
-		CreateTime: timestamppb.New(md.CreationTime),
-		UpdateTime: timestamppb.New(md.LastModifiedTime),
-	})
-	if err != nil {
-		e.logger.Warn("error creating Any struct", "error", err)
-	}
-
 	asset := &v1beta2.Asset{
 		Urn:         tableURN,
 		Name:        t.TableID,
 		Type:        "table",
 		Description: md.Description,
 		Service:     "bigquery",
-		Data:        table,
-		Labels:      md.Labels,
+		// Data:        table,
+		Labels: md.Labels,
+	}
+	attributesData := map[string]interface{}{
+		"full_qualified_name": tableFQN,
+		"dataset":             t.DatasetID,
+		"project":             t.ProjectID,
+		"type":                string(md.Type),
+		"partition_data":      partitionData,
+		"clustering_fields":   clusteringFields,
 	}
 
-	if e.config.BuildViewLineage && (md.Type == bigquery.ViewTable || md.Type == bigquery.MaterializedView) {
+	if md.Type == bigquery.ViewTable || md.Type == bigquery.MaterializedView {
 		query := getViewQuery(md)
-		upstreamResources := getUpstreamResources(query)
-		asset.Lineage = &v1beta2.Lineage{
-			Upstreams: upstreamResources,
+		attributesData["sql"] = query
+		if e.config.BuildViewLineage {
+			upstreamResources := getUpstreamResources(query)
+			asset.Lineage = &v1beta2.Lineage{
+				Upstreams: upstreamResources,
+			}
 		}
 	}
-
+	table, err := anypb.New(&v1beta2.Table{
+		Columns:       e.buildColumns(ctx, md.Schema, md),
+		PreviewFields: previewFields,
+		PreviewRows:   previewRows,
+		Profile:       tableProfile,
+		Attributes:    utils.TryParseMapToProto(attributesData),
+		CreateTime:    timestamppb.New(md.CreationTime),
+		UpdateTime:    timestamppb.New(md.LastModifiedTime),
+	})
+	if err != nil {
+		e.logger.Warn("error creating Any struct", "error", err)
+	}
+	asset.Data = table
 	return asset, nil
 }