fix: Bigquery dataset create table disposition (#4649)

danbaron63 · web-flow · commit 58e03d17448a · 2024-10-21T21:10:14.000+04:00
* fix: adding self to a method which is failing linting in file.py integration tests, added self param to a method that was failing linting and ignoring other issues

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;

* fix: added create_table_disposition check when creating a dataset when get_historical_features is called

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;

* fix: ignoring some mypy linting errors caused by expanding a dict into kwargs in the repo_configuration integration tests

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;

* ignoring typing in auth_permissions_util that would be unreasonable to fix due to length of type required and imports

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;

* fixing method declaration that has no self parameter

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;

* made xdist_group methods static

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;

---------

Signed-off-by: Dan Baron &lt;dan.baron@starlingbank.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -242,6 +242,7 @@ def get_historical_features(
             dataset_project,
             config.offline_store.dataset,
             config.offline_store.location,
+            config.offline_store.table_create_disposition,
         )
 
         entity_schema = _get_entity_schema(
@@ -670,6 +671,7 @@ def _get_table_reference_for_new_entity(
     dataset_project: str,
     dataset_name: str,
     dataset_location: Optional[str],
+    table_create_disposition: str,
 ) -> str:
     """Gets the table_id for the new entity to be uploaded."""
 
@@ -679,8 +681,13 @@ def _get_table_reference_for_new_entity(
 
     try:
         client.get_dataset(dataset.reference)
-    except NotFound:
+    except NotFound as nfe:
         # Only create the dataset if it does not exist
+        if table_create_disposition == "CREATE_NEVER":
+            raise ValueError(
+                f"Dataset {dataset_project}.{dataset_name} does not exist "
+                f"and table_create_disposition is set to {table_create_disposition}."
+            ) from nfe
         client.create_dataset(dataset, exists_ok=True)
 
     table_name = offline_utils.get_temp_entity_table_name()
diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py
@@ -575,9 +575,9 @@ def construct_test_environment(
     }
 
     if not isinstance(offline_creator, RemoteOfflineOidcAuthStoreDataSourceCreator):
-        environment = Environment(**environment_params)
+        environment = Environment(**environment_params)  # type: ignore
     else:
-        environment = OfflineServerPermissionsEnvironment(**environment_params)
+        environment = OfflineServerPermissionsEnvironment(**environment_params)  # type: ignore
     return environment
 
 
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_source_creator.py b/sdk/python/tests/integration/feature_repos/universal/data_source_creator.py
@@ -61,5 +61,6 @@ def create_logged_features_destination(self) -> LoggingDestination:
     def teardown(self):
         raise NotImplementedError
 
+    @staticmethod
     def xdist_groups() -> list[str]:
         return []
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py
@@ -451,6 +451,7 @@ def __init__(self, project_name: str, *args, **kwargs):
         self.server_port: int = 0
         self.proc = None
 
+    @staticmethod
     def xdist_groups() -> list[str]:
         return ["keycloak"]
 
@@ -464,10 +465,10 @@ def setup(self, registry: RegistryConfig):
             entity_key_serialization_version=2,
         )
 
-        repo_path = Path(tempfile.mkdtemp())
-        with open(repo_path / "feature_store.yaml", "w") as outfile:
+        repo_base_path = Path(tempfile.mkdtemp())
+        with open(repo_base_path / "feature_store.yaml", "w") as outfile:
             yaml.dump(config.model_dump(by_alias=True), outfile)
-        repo_path = str(repo_path.resolve())
+        repo_path = str(repo_base_path.resolve())
 
         include_auth_config(
             file_path=f"{repo_path}/feature_store.yaml", auth_config=self.auth_config
@@ -486,7 +487,7 @@ def setup(self, registry: RegistryConfig):
         ]
         self.proc = subprocess.Popen(
             cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL
-        )
+        )  # type: ignore
 
         _time_out_sec: int = 60
         # Wait for server to start
diff --git a/sdk/python/tests/utils/auth_permissions_util.py b/sdk/python/tests/utils/auth_permissions_util.py
@@ -49,7 +49,7 @@ def default_store(
 
     fs = FeatureStore(repo_path=repo_path)
 
-    fs.apply(permissions)
+    fs.apply(permissions)  # type: ignore
 
     return fs
 

Original file line number	Diff line number	Diff line change
`@@ -575,9 +575,9 @@ def construct_test_environment(`
`575`	`575`	`}`
`576`	`576`
`577`	`577`	`if not isinstance(offline_creator, RemoteOfflineOidcAuthStoreDataSourceCreator):`
`578`		`- environment = Environment(**environment_params)`
	`578`	`+ environment = Environment(**environment_params) # type: ignore`
`579`	`579`	`else:`
`580`		`- environment = OfflineServerPermissionsEnvironment(**environment_params)`
	`580`	`+ environment = OfflineServerPermissionsEnvironment(**environment_params) # type: ignore`
`581`	`581`	`return environment`
`582`	`582`
`583`	`583`