generator.dryrun

  1"""Dry Run method to get BigQuery metadata."""
  2
  3import json
  4from enum import Enum
  5from functools import cached_property
  6from typing import Optional
  7from urllib.request import Request, urlopen
  8
  9import google.auth
 10from google.auth.transport.requests import Request as GoogleAuthRequest
 11from google.cloud import bigquery
 12from google.oauth2.id_token import fetch_id_token
 13
 14DRY_RUN_URL = (
 15    "https://us-central1-moz-fx-data-shared-prod.cloudfunctions.net/bigquery-etl-dryrun"
 16)
 17
 18
 19def credentials(auth_req: Optional[GoogleAuthRequest] = None):
 20    """Get GCP credentials."""
 21    auth_req = auth_req or GoogleAuthRequest()
 22    creds, _ = google.auth.default(
 23        scopes=["https://www.googleapis.com/auth/cloud-platform"]
 24    )
 25    creds.refresh(auth_req)
 26    return creds
 27
 28
 29def id_token():
 30    """Get token to authenticate against Cloud Function."""
 31    auth_req = GoogleAuthRequest()
 32    creds = credentials(auth_req)
 33
 34    if hasattr(creds, "id_token"):
 35        # Get token from default credentials for the current environment created via Cloud SDK run
 36        id_token = creds.id_token
 37    else:
 38        # If the environment variable GOOGLE_APPLICATION_CREDENTIALS is set to service account JSON file,
 39        # then ID token is acquired using this service account credentials.
 40        id_token = fetch_id_token(auth_req, DRY_RUN_URL)
 41    return id_token
 42
 43
 44class DryRunError(Exception):
 45    """Exception raised on dry run errors."""
 46
 47    def __init__(self, message, error, use_cloud_function, table_id):
 48        """Initialize DryRunError."""
 49        super().__init__(message)
 50        self.error = error
 51        self.use_cloud_function = use_cloud_function
 52        self.table_id = table_id
 53
 54    def __reduce__(self):
 55        """
 56        Override to ensure that all parameters are being passed when pickling.
 57
 58        Pickling happens when passing exception between processes (e.g. via multiprocessing)
 59        """
 60        return (
 61            self.__class__,
 62            self.args + (self.error, self.use_cloud_function, self.table_id),
 63        )
 64
 65
 66class Errors(Enum):
 67    """DryRun errors that require special handling."""
 68
 69    READ_ONLY = 1
 70    DATE_FILTER_NEEDED = 2
 71    DATE_FILTER_NEEDED_AND_SYNTAX = 3
 72    PERMISSION_DENIED = 4
 73
 74
 75class DryRunContext:
 76    """DryRun builder class."""
 77
 78    def __init__(
 79        self,
 80        use_cloud_function=False,
 81        id_token=None,
 82        credentials=None,
 83        dry_run_url=DRY_RUN_URL,
 84    ):
 85        """Initialize dry run instance."""
 86        self.use_cloud_function = use_cloud_function
 87        self.dry_run_url = dry_run_url
 88        self.id_token = id_token
 89        self.credentials = credentials
 90
 91    def create(
 92        self,
 93        sql=None,
 94        project="moz-fx-data-shared-prod",
 95        dataset=None,
 96        table=None,
 97    ):
 98        """Initialize a DryRun instance."""
 99        return DryRun(
100            use_cloud_function=self.use_cloud_function,
101            id_token=self.id_token,
102            credentials=self.credentials,
103            sql=sql,
104            project=project,
105            dataset=dataset,
106            table=table,
107            dry_run_url=self.dry_run_url,
108        )
109
110
111class DryRun:
112    """Dry run SQL."""
113
114    def __init__(
115        self,
116        use_cloud_function=False,
117        id_token=None,
118        credentials=None,
119        sql=None,
120        project="moz-fx-data-shared-prod",
121        dataset=None,
122        table=None,
123        dry_run_url=DRY_RUN_URL,
124    ):
125        """Initialize dry run instance."""
126        self.sql = sql
127        self.use_cloud_function = use_cloud_function
128        self.project = project
129        self.dataset = dataset
130        self.table = table
131        self.dry_run_url = dry_run_url
132        self.id_token = id_token
133        self.credentials = credentials
134
135    @cached_property
136    def client(self):
137        """Get BigQuery client instance."""
138        return bigquery.Client(credentials=self.credentials)
139
140    @cached_property
141    def dry_run_result(self):
142        """Return the dry run result."""
143        try:
144            if self.use_cloud_function:
145                json_data = {
146                    "query": self.sql or "SELECT 1",
147                    "project": self.project,
148                    "dataset": self.dataset or "telemetry",
149                }
150
151                if self.table:
152                    json_data["table"] = self.table
153
154                r = urlopen(
155                    Request(
156                        self.dry_run_url,
157                        headers={
158                            "Content-Type": "application/json",
159                            "Authorization": f"Bearer {self.id_token}",
160                        },
161                        data=json.dumps(json_data).encode("utf8"),
162                        method="POST",
163                    )
164                )
165                return json.load(r)
166            else:
167                query_schema = None
168                referenced_tables = []
169                table_metadata = None
170
171                if self.sql:
172                    job_config = bigquery.QueryJobConfig(
173                        dry_run=True,
174                        use_query_cache=False,
175                        query_parameters=[
176                            bigquery.ScalarQueryParameter(
177                                "submission_date", "DATE", "2019-01-01"
178                            )
179                        ],
180                    )
181
182                    if self.project:
183                        job_config.connection_properties = [
184                            bigquery.ConnectionProperty(
185                                "dataset_project_id", self.project
186                            )
187                        ]
188
189                    job = self.client.query(self.sql, job_config=job_config)
190                    query_schema = (
191                        job._properties.get("statistics", {})
192                        .get("query", {})
193                        .get("schema", {})
194                    )
195                    referenced_tables = [
196                        ref.to_api_repr() for ref in job.referenced_tables
197                    ]
198
199                if (
200                    self.project is not None
201                    and self.table is not None
202                    and self.dataset is not None
203                ):
204                    table = self.client.get_table(
205                        f"{self.project}.{self.dataset}.{self.table}"
206                    )
207                    table_metadata = {
208                        "tableType": table.table_type,
209                        "friendlyName": table.friendly_name,
210                        "schema": {
211                            "fields": [field.to_api_repr() for field in table.schema]
212                        },
213                    }
214
215                return {
216                    "valid": True,
217                    "referencedTables": referenced_tables,
218                    "schema": query_schema,
219                    "tableMetadata": table_metadata,
220                }
221        except Exception as e:
222            print(f"ERROR {e}")
223            return None
224
225    def get_schema(self):
226        """Return the query schema by dry running the SQL file."""
227        self.validate()
228
229        if (
230            self.dry_run_result
231            and self.dry_run_result["valid"]
232            and "schema" in self.dry_run_result
233        ):
234            return self.dry_run_result["schema"]["fields"]
235
236        return []
237
238    def get_table_schema(self):
239        """Return the schema of the provided table."""
240        self.validate()
241
242        if (
243            self.dry_run_result
244            and self.dry_run_result["valid"]
245            and "tableMetadata" in self.dry_run_result
246        ):
247            return self.dry_run_result["tableMetadata"]["schema"]["fields"]
248
249        return []
250
251    def get_table_metadata(self):
252        """Return table metadata."""
253        self.validate()
254
255        if (
256            self.dry_run_result
257            and self.dry_run_result["valid"]
258            and "tableMetadata" in self.dry_run_result
259        ):
260            return self.dry_run_result["tableMetadata"]
261
262        return {}
263
264    def validate(self):
265        """Dry run the provided SQL file and check if valid."""
266        dry_run_error = DryRunError(
267            "Error when dry running SQL",
268            self.get_error(),
269            self.use_cloud_function,
270            self.table,
271        )
272
273        if self.dry_run_result is None:
274            raise dry_run_error
275
276        if self.dry_run_result["valid"]:
277            return True
278        elif self.get_error() == Errors.READ_ONLY:
279            # We want the dryrun service to only have read permissions, so
280            # we expect CREATE VIEW and CREATE TABLE to throw specific
281            # exceptions.
282            return True
283        elif self.get_error() == Errors.DATE_FILTER_NEEDED:
284            # With strip_dml flag, some queries require a partition filter
285            # (submission_date, submission_timestamp, etc.) to run
286            return True
287        else:
288            print("ERROR\n", self.dry_run_result["errors"])
289            raise dry_run_error
290
291    def errors(self):
292        """Dry run the provided SQL file and return errors."""
293        if self.dry_run_result is None:
294            return []
295        return self.dry_run_result.get("errors", [])
296
297    def get_error(self) -> Optional[Errors]:
298        """Get specific errors for edge case handling."""
299        errors = self.errors()
300        if len(errors) != 1:
301            return None
302
303        error = errors[0]
304        if error and error.get("code") in [400, 403]:
305            error_message = error.get("message", "")
306            if (
307                "does not have bigquery.tables.create permission for dataset"
308                in error_message
309                or "Permission bigquery.tables.create denied" in error_message
310                or "Permission bigquery.datasets.update denied" in error_message
311            ):
312                return Errors.READ_ONLY
313            if "without a filter over column(s)" in error_message:
314                return Errors.DATE_FILTER_NEEDED
315            if (
316                "Syntax error: Expected end of input but got keyword WHERE"
317                in error_message
318            ):
319                return Errors.DATE_FILTER_NEEDED_AND_SYNTAX
320            if (
321                "Permission bigquery.tables.get denied on table" in error_message
322                or "User does not have permission to query table" in error_message
323            ):
324                return Errors.PERMISSION_DENIED
325        return None