feat(bigquery): add better timers around every API call #8626

mayurinehate · 2023-08-14T14:57:02Z

Tested with bigquery connector-tests for accuracy.

Primary changes -

Updated PerfTimer implementation to allow pausing timer
Renamed BigQueryDataDictionary to BigQuerySchemaApi, added timers for the calls in this class.
Refractored core audit log extraction methods into new class BigQueryAuditLogApi, added timers for the calls in this class.

Refractors:

Refractored connection config and related methods in BigQueryConnectionConfig class.
Refractored BigQueryDataDictionary to include bigquery client as instance method and moved related static methods to instance methods.
Refractored lineage code to be solely in BigQueryLineageExtractor.
Refractored and moved all queries/query templates to bigquery_v2/queries.py

Functional changes -

Fix repetitive access to BQ exported audit logs in lineage case
Deduplicates insertId via QUALIFY in BQ exported audit logs in lineage case

Checklist

The PR conforms to DataHub's Contributing Guideline (particularly Commit Message Format)
Links to related issues (if applicable)
Tests for the changes have been added/updated (if applicable)
Docs related to the changes have been added/updated (if applicable). If a new feature has been added a Usage Guide has been added for the same.
For any breaking change/potential downtime/deprecation/big changes an entry has been made in Updating DataHub

…refractor

From this datahub-project@fa0c43c

asikowitz

Had some trouble reviewing this because there's so much copied / moved code, ended up only skimming thru most of the lineage and usage changes. Looks very good to me, my main concerns around assertions in the perf timer. I think we should avoid erroring due to the perf time and instead note something went wrong in the report, if possible

asikowitz · 2023-08-28T21:41:18Z

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

        try:
-            projects = BigQueryDataDictionary.get_projects(conn)
+            projects = self.bigquery_data_dictionary.get_projects()
        except Exception as e:


Might change my mind after seeing the data dictionary, but I think it might make sense to catch this in the data dictionary method rather than here.

not sure if it makes a big difference but I moved it as suggested.

asikowitz · 2023-09-12T21:20:39Z

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

+        bigquery_audit_metadata_query_template: Callable[
+            [
+                str,  # dataset: str
+                bool,  # use_date_sharded_tables: bool
+                Optional[int],  # limit: Optional[int] = None
+            ],
+            str,
+        ],


This is fine here, but in general would like to avoid passing complex functions because it's kinda ugly and hard to extend

asikowitz · 2023-09-12T21:22:26Z

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

+                if rate_limiter:
+                    with rate_limiter:
+                        for entry in query_job:
+                            with current_timer.pause_timer():
+                                yield entry
+                else:
+                    for entry in query_job:
+                        with current_timer.pause_timer():
+                            yield entry


This is pretty awkward. Maybe we can set the rate limiter to an empty context manager or put the decision whether to rate limit inside the rate limiter (would prob have to wrap it then)? Not sure

refractored it slightly. I think it looks better now.

asikowitz · 2023-09-12T21:43:57Z

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

+                if i == 0:
+                    logger.info(
+                        f"Starting log load from GCP Logging for {client.project}"
+                    )


This seems unnecessary -- can gain this info from the next log

removing this

asikowitz · 2023-09-12T22:03:28Z

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

+        self.report = report
+
+    def get_client(self) -> bigquery.Client:
+        assert self.bq_client is not None


This shouldn't be necessary, right?

yep, removing this method altogether.

metadata-ingestion/src/datahub/utilities/perf_timer.py

asikowitz · 2023-09-12T22:57:35Z

metadata-ingestion/src/datahub/utilities/perf_timer.py

+        assert (
+            not self.paused and not self.end_time
+        ), "Can not pause a paused/stopped timer"
+        assert (
+            self.start_time is not None
+        ), "Can not pause a timer that hasn't started. Did you forget to start the timer ?"


For safety, since we're not great about wrapping everything with try / excepts, what do you think about a logger warning or error here instead? I would rather complete ingestion with faulty timers than have an uncaught exception. Or perhaps we can store an error state in this object, so it can be seen in the report and we know not to trust the time

Right. Let me add warning logs in this case, and also add the error state indicator for reported time in such case.

asikowitz · 2023-09-12T22:59:24Z

metadata-ingestion/src/datahub/utilities/perf_timer.py

+        if self.paused:  # Entering paused timer context, NO OP
+            pass
+        else:
+            self.start()


Suggested change

if self.paused: # Entering paused timer context, NO OP

pass

else:

self.start()

if not self.paused:

self.start()

asikowitz · 2023-09-12T23:00:02Z

metadata-ingestion/src/datahub/utilities/perf_timer.py

+        if self.paused:
+            self.paused = False
+
+    def pause_timer(self) -> "PerfTimer":


Thoughts on renaming this to just pause(). Generally it's clear that the object being paused is a timer

makes sense.

asikowitz · 2023-09-12T23:02:42Z

metadata-ingestion/tests/unit/utilities/test_perf_timer.py

+        time.sleep(1)
+        assert round(timer.elapsed_seconds()) == 1
+
+    assert round(timer.elapsed_seconds()) == 1


Suggested change

assert round(timer.elapsed_seconds()) == 1

assert pytest.approx(timer.elapsed_seconds()) == 1

Can do the same for others

asikowitz · 2023-09-13T18:11:40Z

Looks like a legitimate failure on FAILED tests/unit/test_bigquery_source.py::test_get_projects_list_failure because of the project list change

mayurinehate · 2023-09-15T06:45:32Z

Looks like a legitimate failure on FAILED tests/unit/test_bigquery_source.py::test_get_projects_list_failure because of the project list change

@asikowitz the test is fixed.

Follow up on datahub-project#8626.

mayurinehate added 5 commits August 9, 2023 20:45

feat(bigquery): add better timers around every API call

b84601c

wip, timers not added for unused methods - remove these ?

95bbcbe

refractor in lineage.py

76ddc3f

report composition vs inheritance

e57f134

more refractor and fixes

38b18bb

github-actions bot added the ingestion PR or Issue related to the ingestion of metadata label Aug 14, 2023

Merge branch 'master' into bq_timers_api_refractor

8d049a7

vercel bot deployed to Preview August 14, 2023 15:32 View deployment

fix lint, tests

31a3be8

vercel bot deployed to Preview August 14, 2023 15:59 View deployment

mayurinehate mentioned this pull request Aug 17, 2023

fix(ingest): stateful redundant run skip handler #8467

Merged

5 tasks

asikowitz self-requested a review August 21, 2023 17:16

mayurinehate added 2 commits August 22, 2023 12:16

revert rename of bigquery_schema.py to bigquery_schema_api.py

eaa72a3

Merge remote-tracking branch 'datahub-oss/master' into bq_timers_api_…

ac2ab3b

…refractor

vercel bot deployed to Preview August 22, 2023 08:51 View deployment

Merge branch 'master' into bq_timers_api_refractor

94077db

vercel bot deployed to Preview August 28, 2023 06:16 View deployment

mayurinehate added 2 commits August 28, 2023 18:58

Merge remote-tracking branch 'datahub-oss/master' into bq_timers_api_…

e500275

…refractor

move stateful check inside lineage module

1b3d5b5

vercel bot deployed to Preview August 28, 2023 14:45 View deployment

mayurinehate requested a review from hsheth2 August 28, 2023 14:47

Merge branch 'master' into bq_timers_api_refractor

c77a9ab

vercel bot deployed to Preview September 5, 2023 12:27 View deployment

merge related changes

6a2a3d4

From this datahub-project@fa0c43c

mayurinehate force-pushed the bq_timers_api_refractor branch from b3236fc to 6a2a3d4 Compare September 5, 2023 12:29

vercel bot deployed to Preview September 5, 2023 12:56 View deployment

hsheth2 assigned asikowitz Sep 6, 2023

Merge branch 'master' into bq_timers_api_refractor

32fedfc

vercel bot deployed to Preview September 8, 2023 04:58 View deployment

Merge branch 'master' into bq_timers_api_refractor

9dca7e5

vercel bot deployed to Preview September 11, 2023 06:42 View deployment

asikowitz reviewed Sep 12, 2023

View reviewed changes

mayurinehate and others added 2 commits September 13, 2023 17:43

address review comments

79f84ba

Merge branch 'master' into bq_timers_api_refractor

25f4f0b

vercel bot deployed to Preview September 13, 2023 12:46 View deployment

asikowitz approved these changes Sep 13, 2023

View reviewed changes

fix tests

0ba4efc

vercel bot deployed to Preview September 14, 2023 06:04 View deployment

Merge branch 'master' into bq_timers_api_refractor

ea2c4a1

vercel bot deployed to Preview September 15, 2023 09:12 View deployment

hsheth2 merged commit cdb9f5b into datahub-project:master Sep 15, 2023
57 checks passed

hsheth2 added a commit to hsheth2/datahub that referenced this pull request Sep 20, 2023

fix(ingest/bigquery): show report in output

dac739a

Follow up on datahub-project#8626.

hsheth2 mentioned this pull request Sep 20, 2023

fix(ingest/bigquery): show report in output #8867

Merged

5 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(bigquery): add better timers around every API call #8626

feat(bigquery): add better timers around every API call #8626

mayurinehate commented Aug 14, 2023 •

edited

Loading

asikowitz left a comment

asikowitz Aug 28, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz Sep 12, 2023

mayurinehate Sep 13, 2023

asikowitz commented Sep 13, 2023

mayurinehate commented Sep 15, 2023 •

edited

Loading

	assert round(timer.elapsed_seconds()) == 1
	assert pytest.approx(timer.elapsed_seconds()) == 1

feat(bigquery): add better timers around every API call #8626

feat(bigquery): add better timers around every API call #8626

Conversation

mayurinehate commented Aug 14, 2023 • edited Loading

Checklist

asikowitz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

asikowitz commented Sep 13, 2023

mayurinehate commented Sep 15, 2023 • edited Loading

mayurinehate commented Aug 14, 2023 •

edited

Loading

mayurinehate commented Sep 15, 2023 •

edited

Loading