Building a Self-Hosted BigQuery Data Catalog on GKE Autopilot

1from google.cloud import bigquery
2from google.cloud import datacatalog_lineage_v1
3import vertexai
4from vertexai.generative_models import GenerativeModel
5
6bq_client = bigquery.Client()
7lineage_client = datacatalog_lineage_v1.LineageClient()
8model = GenerativeModel("gemini-2.5-flash")
9
10async def sync_dataset(project_id: str, dataset_id: str, db: AsyncSession):
11    dataset = bq_client.get_dataset(f"{project_id}.{dataset_id}")
12
13    for table_ref in bq_client.list_tables(dataset):
14        table = bq_client.get_table(table_ref)
15
16        # Upsert table metadata
17        await upsert_table(db, {
18            "project_id": project_id,
19            "dataset_id": dataset_id,
20            "table_id": table.table_id,
21            "row_count": table.num_rows,
22            "size_bytes": table.num_bytes,
23            "last_modified": table.modified,
24        })
25
26        for field in table.schema:
27            existing = await get_column(db, project_id, dataset_id, table.table_id, field.name)
28
29            # Generate AI description only for new/undocumented columns
30            if not existing or not existing.ai_description:
31                description = await generate_column_description(
32                    model, table, field
33                )
34                await upsert_column(db, {
35                    "column_name": field.name,
36                    "data_type": field.field_type,
37                    "ai_description": description,
38                    "is_nullable": field.is_nullable,
39                })

1async def generate_column_description(
2    model: GenerativeModel,
3    table: bigquery.Table,
4    field: bigquery.SchemaField,
5) -> str:
6    # Get sample values for context
7    sample_query = f"""
8        SELECT DISTINCT {field.name}
9        FROM `{table.full_table_id.replace(':', '.')}`
10        WHERE {field.name} IS NOT NULL
11        LIMIT 5
12    """
13    sample_rows = list(bq_client.query(sample_query))
14    sample_values = [str(row[field.name]) for row in sample_rows]
15
16    prompt = f"""You are a data dictionary expert. Generate a concise, accurate
17description for this BigQuery column.
18
19Table: {table.table_id} (in dataset {table.dataset_id})
20Table description: {table.description or 'Not documented'}
21Column name: {field.name}
22Data type: {field.field_type}
23Sample values: {', '.join(sample_values)}
24
25Write a 1-2 sentence description explaining what this column represents,
26what the values mean, and any business context implied by the name or values.
27Be specific — avoid generic descriptions like 'a field containing data'."""
28
29    response = model.generate_content(prompt)
30    return response.text.strip()

1# Create GCP service account
2gcloud iam service-accounts create ds-catalog-sa   --display-name="DS Data Catalog Service Account"
3
4# Grant required permissions
5gcloud projects add-iam-policy-binding $PROJECT_ID   --member="serviceAccount:ds-catalog-sa@$PROJECT_ID.iam.gserviceaccount.com"   --role="roles/bigquery.metadataViewer"
6
7# Bind K8s SA to GCP SA (the Workload Identity link)
8gcloud iam service-accounts add-iam-policy-binding   ds-catalog-sa@$PROJECT_ID.iam.gserviceaccount.com   --role="roles/iam.workloadIdentityUser"   --member="serviceAccount:$PROJECT_ID.svc.id.goog[default/ds-catalog-ksa]"
9
10# Annotate K8s service account
11kubectl annotate serviceaccount ds-catalog-ksa   iam.gke.io/gcp-service-account=ds-catalog-sa@$PROJECT_ID.iam.gserviceaccount.com

1def compute_quality_score(table_stats: dict) -> float:
2    """Compute a 0–100 quality score for a BigQuery table."""
3    score = 100.0
4
5    # Freshness penalty: -5 points per day since last modification
6    days_stale = (datetime.now() - table_stats["last_modified"]).days
7    score -= min(days_stale * 5, 40)  # Max -40 for freshness
8
9    # Documentation penalty: -20 if no table description
10    if not table_stats.get("description"):
11        score -= 20
12
13    # Column coverage: -1 point per undocumented column (max -20)
14    total_cols = table_stats["column_count"]
15    documented_cols = table_stats["documented_column_count"]
16    if total_cols > 0:
17        undocumented_pct = (total_cols - documented_cols) / total_cols
18        score -= min(undocumented_pct * 20, 20)
19
20    # Row count sanity: -10 if table is empty
21    if table_stats.get("row_count", 0) == 0:
22        score -= 10
23
24    return max(0, round(score, 1))