Add limit to execute signature for consistency with dbt-core

.changes/unreleased/Fixes-20230509-222705.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: Update signature for execute method
+time: 2023-05-09T22:27:05.976611-07:00
+custom:
+  Author: aranke
+  Issue: ' '

dbt/adapters/bigquery/connections.py

@@ -410,7 +410,7 @@ def get_table_from_response(cls, resp):
         column_names = [field.name for field in resp.schema]
         return agate_helper.table_from_data_flat(resp, column_names)
 
-    def raw_execute(self, sql, fetch=False, *, use_legacy_sql=False):
+    def raw_execute(self, sql, use_legacy_sql=False, limit: Optional[int] = None):
         conn = self.get_thread_connection()
         client = conn.handle
 
@@ -420,8 +420,8 @@ def raw_execute(self, sql, fetch=False, *, use_legacy_sql=False):
             and self.profile.query_comment
             and self.profile.query_comment.job_label
         ):
-            query_comment = self.query_header.comment.query_comment
-            labels = self._labels_from_query_comment(query_comment)
+            query_comment = self.profile.query_comment
+            labels = self._labels_from_query_comment(query_comment.comment)
         else:
             labels = {}
 
@@ -450,18 +450,19 @@ def fn():
                 job_params,
                 job_creation_timeout=job_creation_timeout,
                 job_execution_timeout=job_execution_timeout,
+                limit=limit,
             )
 
         query_job, iterator = self._retry_and_handle(msg=sql, conn=conn, fn=fn)
 
         return query_job, iterator
 
     def execute(
-        self, sql, auto_begin=False, fetch=None
+        self, sql, auto_begin=False, fetch=None, limit: Optional[int] = None
     ) -> Tuple[BigQueryAdapterResponse, agate.Table]:
         sql = self._add_query_comment(sql)
         # auto_begin is ignored on bigquery, and only included for consistency
-        query_job, iterator = self.raw_execute(sql, fetch=fetch)
+        query_job, iterator = self.raw_execute(sql, limit=limit)
 
         if fetch:
             table = self.get_table_from_response(iterator)
@@ -550,7 +551,7 @@ def standard_to_legacy(table):
 
         sql = self._add_query_comment(legacy_sql)
         # auto_begin is ignored on bigquery, and only included for consistency
-        _, iterator = self.raw_execute(sql, fetch="fetch_result", use_legacy_sql=True)
+        _, iterator = self.raw_execute(sql, use_legacy_sql=True)
         return self.get_table_from_response(iterator)
 
     def copy_bq_table(self, source, destination, write_disposition):
@@ -644,12 +645,13 @@ def _query_and_results(
         job_params,
         job_creation_timeout=None,
         job_execution_timeout=None,
+        limit: Optional[int] = None,
     ):
         """Query the client and wait for results."""
         # Cannot reuse job_config if destination is set and ddl is used
         job_config = google.cloud.bigquery.QueryJobConfig(**job_params)
         query_job = client.query(query=sql, job_config=job_config, timeout=job_creation_timeout)
-        iterator = query_job.result(timeout=job_execution_timeout)
+        iterator = query_job.result(max_results=limit, timeout=job_execution_timeout)
 
         return query_job, iterator