Speed up view deploys

bigquery_etl/cli/deploy.py

@@ -42,6 +42,7 @@
 from bigquery_etl.routine.parse_routine import ROUTINE_FILES, RawRoutine
 from bigquery_etl.schema import SCHEMA_FILE, Schema
 from bigquery_etl.util import extract_from_query_path
+from bigquery_etl.util.client_queue import get_client
 from bigquery_etl.util.common import block_coding_agents, get_bqetl_project_root, render
 from bigquery_etl.util.parallel_topological_sorter import ParallelTopologicalSorter
 from bigquery_etl.util.target import (
@@ -1271,7 +1272,7 @@ def _deploy_table_artifact(file_path: Path, options: dict):
 
         # validate schema matches query if not using --table-force
         if not options["table_force"] and str(file_path).endswith(".sql"):
-            client = bigquery.Client(credentials=options["credentials"])
+            client = get_client(credentials=options["credentials"])
             try:
                 table_name = file_path.parent.name
                 dataset_name = file_path.parent.parent.name
@@ -1309,6 +1310,7 @@ def _deploy_table_artifact(file_path: Path, options: dict):
         credentials=options["credentials"],
         id_token=options["id_token"],
         isolated=options.get("isolated", False),
+        client=get_client(credentials=options["credentials"]),
     )
 
 
@@ -1317,7 +1319,8 @@ def _deploy_view_artifact(file_path: Path, options: dict):
     id_token = options.get("id_token")
     view = View.from_file(file_path, id_token=id_token)
 
-    # Add managed label if requested
+    # Add managed label if requested. This has to happen before the
+    # has_changes check below
     if options.get("view_add_managed_label", False):
         view.labels["managed"] = ""
 
@@ -1326,11 +1329,27 @@ def _deploy_view_artifact(file_path: Path, options: dict):
             raise FailedDeployException(f"View validation failed for {file_path}")
         return
 
-    client = bigquery.Client(credentials=options["credentials"])
+    client = get_client(credentials=options["credentials"])
+    target_project = options.get("view_target_project")
+
+    if not options["view_force"]:
+        # Read outside the try so a missing option raises a KeyError
+        use_cloud_function = options["use_cloud_function"]
+        try:
+            changed = view.has_changes(
+                target_project,
+                client=client,
+                use_cloud_function=use_cloud_function,
+            )
+        except Exception as e:
+            # The check is an optimization so a failure here just republishes
+            log.warning(f"Could not check {file_path} for changes, publishing: {e}")
+            changed = True
+        if not changed:
+            raise SkippedDeployException("no changes")
 
     success = view.publish(
-        target_project=options.get("view_target_project"),
-        force=options["view_force"],
+        target_project=target_project,
         client=client,
         dry_run=False,
     )

bigquery_etl/cli/view.py

@@ -11,7 +11,6 @@
 from traceback import print_exc
 
 import rich_click as click
-from google.cloud import bigquery
 
 from ..cli.utils import (
     parallelism_option,
@@ -23,8 +22,7 @@
 from ..config import ConfigLoader
 from ..dryrun import DryRun, get_credentials, get_id_token
 from ..metadata.parse_metadata import METADATA_FILE, Metadata
-from ..util.bigquery_id import sql_table_id
-from ..util.client_queue import ClientQueue
+from ..util.client_queue import ClientQueue, get_client
 from ..util.common import block_coding_agents
 from ..util.parallel_topological_sorter import ParallelTopologicalSorter
 from ..view import View, broken_views
@@ -263,7 +261,7 @@ def publish(
 
 
 def _view_has_changes(target_project, credentials, view):
-    return view.has_changes(target_project, credentials)
+    return view.has_changes(target_project, client=get_client(credentials))
 
 
 def _publish_view_callback(
@@ -276,7 +274,7 @@ def _publish_view_callback(
     results,
 ):
     try:
-        client = bigquery.Client(credentials=credentials)
+        client = get_client(credentials)
         success = views_by_id[view_id].publish(target_project, dry_run, client)
         results[view_id] = success if success is not None else True
     except Exception:
@@ -375,6 +373,14 @@ def _collect_views(
     is_flag=True,
     help="Only clean views with labels: {authorized: true} in metadata.yaml",
 )
+@click.option(
+    "--region",
+    default="us",
+    help=(
+        "BigQuery region whose INFORMATION_SCHEMA is queried to find managed views."
+        " Views in datasets outside this region are not cleaned."
+    ),
+)
 def clean(
     name,
     sql_dir,
@@ -386,6 +392,7 @@ def clean(
     user_facing_only,
     skip_authorized,
     authorized_only,
+    region,
 ):
     """Clean managed views."""
     # set log level
@@ -416,39 +423,17 @@ def clean(
 
     client_q = ClientQueue([project_id], parallelism)
     with client_q.client() as client:
-        datasets = [
-            dataset
-            for dataset in client.list_datasets(target_project)
-            if not user_facing_only
-            or not dataset.dataset_id.endswith(
-                tuple(
-                    ConfigLoader.get(
-                        "default", "non_user_facing_dataset_suffixes", fallback=[]
-                    )
-                )
-            )
-        ]
+        managed_view_ids = _list_managed_views(
+            client,
+            target_project or project_id,
+            region,
+            name,
+            skip_authorized,
+            authorized_only,
+            user_facing_only,
+        )
 
     with ThreadPool(parallelism) as p:
-        managed_view_ids = {
-            view
-            for views in p.starmap(
-                client_q.with_client,
-                (
-                    (
-                        _list_managed_views,
-                        dataset,
-                        name,
-                        skip_authorized,
-                        authorized_only,
-                    )
-                    for dataset in datasets
-                ),
-                chunksize=1,
-            )
-            for view in views
-        }
-
         remove_view_ids = sorted(managed_view_ids - expected_view_ids)
         p.starmap(
             client_q.with_client,
@@ -458,16 +443,24 @@ def clean(
 
 
 def _list_managed_views(
-    client, dataset, pattern, skip_authorized, authorized_only=False
+    client,
+    project,
+    region,
+    pattern,
+    skip_authorized,
+    authorized_only=False,
+    user_facing_only=False,
 ):
+    """Return the ids of managed views in the given region of the project."""
     query = f"""
       SELECT
         table_catalog || "." || table_schema || "." || table_name AS table_id,
+        table_schema,
         CONTAINS_SUBSTR(option_value, 'STRUCT("authorized", "")') AS is_authorized
       FROM
-        `{dataset.project}.{dataset.dataset_id}.INFORMATION_SCHEMA.VIEWS`
+        `{project}.region-{region}.INFORMATION_SCHEMA.VIEWS`
       INNER JOIN
-        `{dataset.project}.{dataset.dataset_id}.INFORMATION_SCHEMA.TABLE_OPTIONS`
+        `{project}.region-{region}.INFORMATION_SCHEMA.TABLE_OPTIONS`
       USING
         (table_catalog,
          table_schema,
@@ -477,16 +470,24 @@ def _list_managed_views(
       AND CONTAINS_SUBSTR(option_value, 'STRUCT("managed", "")')
     """
 
+    non_user_facing_suffixes = tuple(
+        ConfigLoader.get("default", "non_user_facing_dataset_suffixes", fallback=[])
+    )
+
     # running a query against information schema instead of using the API to list tables is much faster
     job = client.query(query)
     result = list(job.result())
-    return [
+    return {
         row.table_id
         for row in result
-        if (pattern is None or fnmatchcase(sql_table_id(row.table_id), f"*{pattern}"))
+        if (pattern is None or fnmatchcase(row.table_id, f"*{pattern}"))
         and (not skip_authorized or not row.is_authorized)
         and (not authorized_only or row.is_authorized)
-    ]
+        and (
+            not user_facing_only
+            or not row.table_schema.endswith(non_user_facing_suffixes)
+        )
+    }
 
 
 def _remove_view(client, view_id, dry_run):

bigquery_etl/deploy.py

@@ -41,6 +41,7 @@ def deploy_table(
     credentials=None,
     id_token=None,
     isolated: bool = False,
+    client: Optional[bigquery.Client] = None,
 ) -> None:
     """Deploy a query to a destination."""
     # Under --isolated, schema.yaml is authoritative and the table must
@@ -94,7 +95,7 @@ def deploy_table(
     except Exception as e:  # TODO: Raise/catch more specific exception
         raise SkippedDeployException(f"Schema missing for {artifact_file}.") from e
 
-    client = bigquery.Client(credentials=credentials)
+    client = client or bigquery.Client(credentials=credentials)
     # Skip the query-vs-schema dryrun check only in the isolated + skip-listed
     # case (rewritten refs in the target tree won't resolve in prod). Tying
     # this to isolated rather than dryrun_skipped alone keeps the bypass

bigquery_etl/schema/__init__.py

@@ -135,14 +135,23 @@ def get_field(self, field_path: str) -> dict:
         else:
             raise KeyError(f"Field not found: {field_path}")
 
-    def deploy(self, destination_table: str) -> bigquery.Table:
-        """Deploy the schema to BigQuery named after destination_table."""
-        client = bigquery.Client()
+    def deploy(
+        self,
+        destination_table: str,
+        client: Optional[bigquery.Client] = None,
+        table: Optional[bigquery.Table] = None,
+    ) -> bigquery.Table:
+        """Deploy the schema to BigQuery named after destination_table.
+
+        Pass `table` when the caller already fetched it to skip a redundant get_table call.
+        """
+        client = client or bigquery.Client()
         bigquery_schema = self.to_bigquery_schema()
 
         try:
-            # destination table already exists, update schema
-            table = client.get_table(destination_table)
+            if table is None:
+                # destination table already exists, update schema
+                table = client.get_table(destination_table)
             table.schema = bigquery_schema
             return client.update_table(table, ["schema"])
         except NotFound:

bigquery_etl/util/client_queue.py

@@ -1,13 +1,41 @@
-"""Queue for balancing jobs across billing projects."""
+"""Wrappers around a BigQuery client.
+
+`ClientQueue` shares a fixed set of clients across threads in one process,
+balancing work over billing projects.
+
+`get_client` is for worker processes, which can't share a client since the HTTP session isn't
+fork safe and doesn't pickle, so each caches its own.
+"""
 
 import asyncio
+import os
 from contextlib import contextmanager
+from functools import cache
 from queue import Queue
 
 from google.cloud import bigquery
 from requests import adapters
 
 
+@cache
+def _client(pid, credentials, project):
+    """Return a cached client.
+
+    Keyed by (pid, project, id(credentials)). Under spawn (macOS) a child process re-imports this
+    module and starts empty, but under fork (Linux) it inherits the parent's entries.
+    Keying on pid makes the child build its own client.
+    """
+    return bigquery.Client(credentials=credentials, project=project)
+
+
+def get_client(credentials=None, project=None) -> bigquery.Client:
+    """Return a BigQuery client, reused for the life of this process.
+
+    Reusing a client saves a few hundred milliseconds of credential setup during construction.
+    """
+    return _client(os.getpid(), credentials, project)
+
+
 class ClientQueue:
     """Queue for balancing jobs across billing projects.