@@ -119,6 +119,7 @@ def create(
119119 labels : Optional [Dict [str , str ]] = None ,
120120 encryption_spec_key_name : Optional [str ] = None ,
121121 sync : bool = True ,
122+ create_request_timeout : Optional [float ] = None ,
122123 ) -> "_Dataset" :
123124 """Creates a new dataset and optionally imports data into dataset when
124125 source and import_schema_uri are passed.
@@ -203,6 +204,8 @@ def create(
203204 Whether to execute this method synchronously. If False, this method
204205 will be executed in concurrent Future and any downstream object will
205206 be immediately returned and synced when the Future has completed.
207+ create_request_timeout (float):
208+ Optional. The timeout for the create request in seconds.
206209
207210 Returns:
208211 dataset (Dataset):
@@ -240,6 +243,7 @@ def create(
240243 encryption_spec_key_name = encryption_spec_key_name
241244 ),
242245 sync = sync ,
246+ create_request_timeout = create_request_timeout ,
243247 )
244248
245249 @classmethod
@@ -258,6 +262,8 @@ def _create_and_import(
258262 labels : Optional [Dict [str , str ]] = None ,
259263 encryption_spec : Optional [gca_encryption_spec .EncryptionSpec ] = None ,
260264 sync : bool = True ,
265+ create_request_timeout : Optional [float ] = None ,
266+ import_request_timeout : Optional [float ] = None ,
261267 ) -> "_Dataset" :
262268 """Creates a new dataset and optionally imports data into dataset when
263269 source and import_schema_uri are passed.
@@ -313,6 +319,10 @@ def _create_and_import(
313319 Whether to execute this method synchronously. If False, this method
314320 will be executed in concurrent Future and any downstream object will
315321 be immediately returned and synced when the Future has completed.
322+ create_request_timeout (float):
323+ Optional. The timeout for the create request in seconds.
324+ import_request_timeout (float):
325+ Optional. The timeout for the import request in seconds.
316326
317327 Returns:
318328 dataset (Dataset):
@@ -328,6 +338,7 @@ def _create_and_import(
328338 request_metadata = request_metadata ,
329339 labels = labels ,
330340 encryption_spec = encryption_spec ,
341+ create_request_timeout = create_request_timeout ,
331342 )
332343
333344 _LOGGER .log_create_with_lro (cls , create_dataset_lro )
@@ -345,18 +356,26 @@ def _create_and_import(
345356
346357 # Import if import datasource is DatasourceImportable
347358 if isinstance (datasource , _datasources .DatasourceImportable ):
348- dataset_obj ._import_and_wait (datasource )
359+ dataset_obj ._import_and_wait (
360+ datasource , import_request_timeout = import_request_timeout
361+ )
349362
350363 return dataset_obj
351364
352- def _import_and_wait (self , datasource ):
365+ def _import_and_wait (
366+ self ,
367+ datasource ,
368+ import_request_timeout : Optional [float ] = None ,
369+ ):
353370 _LOGGER .log_action_start_against_resource (
354371 "Importing" ,
355372 "data" ,
356373 self ,
357374 )
358375
359- import_lro = self ._import (datasource = datasource )
376+ import_lro = self ._import (
377+ datasource = datasource , import_request_timeout = import_request_timeout
378+ )
360379
361380 _LOGGER .log_action_started_against_resource_with_lro (
362381 "Import" , "data" , self .__class__ , import_lro
@@ -377,6 +396,7 @@ def _create(
377396 request_metadata : Sequence [Tuple [str , str ]] = (),
378397 labels : Optional [Dict [str , str ]] = None ,
379398 encryption_spec : Optional [gca_encryption_spec .EncryptionSpec ] = None ,
399+ create_request_timeout : Optional [float ] = None ,
380400 ) -> operation .Operation :
381401 """Creates a new managed dataset by directly calling API client.
382402
@@ -419,6 +439,8 @@ def _create(
419439 resource is created.
420440
421441 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
442+ create_request_timeout (float):
443+ Optional. The timeout for the create request in seconds.
422444 Returns:
423445 operation (Operation):
424446 An object representing a long-running operation.
@@ -433,25 +455,33 @@ def _create(
433455 )
434456
435457 return api_client .create_dataset (
436- parent = parent , dataset = gapic_dataset , metadata = request_metadata
458+ parent = parent ,
459+ dataset = gapic_dataset ,
460+ metadata = request_metadata ,
461+ timeout = create_request_timeout ,
437462 )
438463
439464 def _import (
440465 self ,
441466 datasource : _datasources .DatasourceImportable ,
467+ import_request_timeout : Optional [float ] = None ,
442468 ) -> operation .Operation :
443469 """Imports data into managed dataset by directly calling API client.
444470
445471 Args:
446472 datasource (_datasources.DatasourceImportable):
447473 Required. Datasource for importing data to an existing dataset for Vertex AI.
474+ import_request_timeout (float):
475+ Optional. The timeout for the import request in seconds.
448476
449477 Returns:
450478 operation (Operation):
451479 An object representing a long-running operation.
452480 """
453481 return self .api_client .import_data (
454- name = self .resource_name , import_configs = [datasource .import_data_config ]
482+ name = self .resource_name ,
483+ import_configs = [datasource .import_data_config ],
484+ timeout = import_request_timeout ,
455485 )
456486
457487 @base .optional_sync (return_input_arg = "self" )
@@ -461,6 +491,7 @@ def import_data(
461491 import_schema_uri : str ,
462492 data_item_labels : Optional [Dict ] = None ,
463493 sync : bool = True ,
494+ import_request_timeout : Optional [float ] = None ,
464495 ) -> "_Dataset" :
465496 """Upload data to existing managed dataset.
466497
@@ -498,6 +529,8 @@ def import_data(
498529 Whether to execute this method synchronously. If False, this method
499530 will be executed in concurrent Future and any downstream object will
500531 be immediately returned and synced when the Future has completed.
532+ import_request_timeout (float):
533+ Optional. The timeout for the import request in seconds.
501534
502535 Returns:
503536 dataset (Dataset):
@@ -510,7 +543,9 @@ def import_data(
510543 data_item_labels = data_item_labels ,
511544 )
512545
513- self ._import_and_wait (datasource = datasource )
546+ self ._import_and_wait (
547+ datasource = datasource , import_request_timeout = import_request_timeout
548+ )
514549 return self
515550
516551 # TODO(b/174751568) add optional sync support
0 commit comments