diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py index 3ee17174619..ed635ae2145 100644 --- 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 index 73f99fb7c28..c688a848362 100644 --- 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 index aa46160358f..513a94ee210 100644 --- 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 index d8b75aca248..35325c2737e 100644 --- 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 index b8ca7355e98..49ddd1b530d 100644 --- 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