Migrate Pydantic to v2

sdk/python/feast/infra/contrib/spark_kafka_processor.py

@@ -22,7 +22,6 @@ class SparkProcessorConfig(ProcessorConfig):
     processing_time: str
     query_timeout: int
 
-
 class SparkKafkaProcessor(StreamProcessor):
     spark: SparkSession
     format: str

sdk/python/feast/infra/feature_servers/aws_lambda/config.py

@@ -1,5 +1,6 @@
+from typing_extensions import Literal
+
 from pydantic import StrictBool, StrictStr
-from pydantic.typing import Literal
 
 from feast.infra.feature_servers.base_config import BaseFeatureServerConfig
 

sdk/python/feast/infra/feature_servers/base_config.py

@@ -30,5 +30,5 @@ class BaseFeatureServerConfig(FeastConfigBaseModel):
     enabled: StrictBool = False
     """Whether the feature server should be launched."""
 
-    feature_logging: Optional[FeatureLoggingConfig]
+    feature_logging: Optional[FeatureLoggingConfig] = None
     """ Feature logging configuration """

sdk/python/feast/infra/feature_servers/gcp_cloudrun/config.py

@@ -1,5 +1,6 @@
+from typing_extensions import Literal
+
 from pydantic import StrictBool
-from pydantic.typing import Literal
 
 from feast.infra.feature_servers.base_config import BaseFeatureServerConfig
 

sdk/python/feast/infra/feature_servers/local_process/config.py

@@ -1,4 +1,4 @@
-from pydantic.typing import Literal
+from typing_extensions import Literal
 
 from feast.infra.feature_servers.base_config import BaseFeatureServerConfig
 

sdk/python/feast/infra/materialization/snowflake_engine.py

@@ -7,7 +7,7 @@
 import click
 import pandas as pd
 from colorama import Fore, Style
-from pydantic import Field, StrictStr
+from pydantic import ConfigDict, Field, StrictStr
 from pytz import utc
 from tqdm import tqdm
 
@@ -72,9 +72,7 @@ class SnowflakeMaterializationEngineConfig(FeastConfigBaseModel):
 
     schema_: Optional[str] = Field("PUBLIC", alias="schema")
     """ Snowflake schema name """
-
-    class Config:
-        allow_population_by_field_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 @dataclass

sdk/python/feast/infra/offline_stores/bigquery.py

@@ -14,13 +14,13 @@
     Tuple,
     Union,
 )
+from typing_extensions import Literal
 
 import numpy as np
 import pandas as pd
 import pyarrow
 import pyarrow.parquet
-from pydantic import ConstrainedStr, StrictStr, validator
-from pydantic.typing import Literal
+from pydantic import StringConstraints, StrictStr, validator
 from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
 
 from feast import flags_helper
@@ -72,7 +72,7 @@ def get_http_client_info():
     return http_client_info.ClientInfo(user_agent=get_user_agent())
 
 
-class BigQueryTableCreateDisposition(ConstrainedStr):
+class BigQueryTableCreateDisposition(StringConstraints):
     """Custom constraint for table_create_disposition. To understand more, see:
     https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationLoad.FIELDS.create_disposition"""
 
@@ -105,6 +105,8 @@ class BigQueryOfflineStoreConfig(FeastConfigBaseModel):
     table_create_disposition: Optional[BigQueryTableCreateDisposition] = None
     """ (optional) Specifies whether the job is allowed to create new tables. The default value is CREATE_IF_NEEDED."""
 
+    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
     @validator("billing_project_id")
     def project_id_exists(cls, v, values, **kwargs):
         if v and not values["project_id"]:
@@ -515,7 +517,7 @@ def to_bigquery(
                 temp_dest_table = f"{tmp_dest['projectId']}.{tmp_dest['datasetId']}.{tmp_dest['tableId']}"
 
                 # persist temp table
-                sql = f"CREATE TABLE `{dest}` AS SELECT * FROM {temp_dest_table}"
+                sql = f"CREATE TABLE `{dest}` AS SELECT * FROM `{temp_dest_table}`"
                 self._execute_query(sql, timeout=timeout)
 
             print(f"Done writing to '{dest}'.")

sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/athena.py

@@ -13,12 +13,12 @@
     Union,
 )
 
+from typing_extensions import Literal
 import numpy as np
 import pandas as pd
 import pyarrow
 import pyarrow as pa
 from pydantic import StrictStr
-from pydantic.typing import Literal
 from pytz import utc
 
 from feast import OnDemandFeatureView

sdk/python/feast/infra/offline_stores/contrib/mssql_offline_store/mssql.py

@@ -3,15 +3,14 @@
 import warnings
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union, Literal
 
 import numpy as np
 import pandas
 import pyarrow
 import pyarrow as pa
 import sqlalchemy
 from pydantic.types import StrictStr
-from pydantic.typing import Literal
 from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import sessionmaker

sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py

@@ -14,12 +14,12 @@
     Union,
 )
 
+from typing_extensions import Literal
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 from jinja2 import BaseLoader, Environment
 from psycopg2 import sql
-from pydantic.typing import Literal
 from pytz import utc
 
 from feast.data_source import DataSource

sdk/python/feast/infra/offline_stores/file.py

@@ -3,14 +3,14 @@
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Callable, List, Optional, Tuple, Union
+from typing_extensions import Literal
 
 import dask.dataframe as dd
 import pandas as pd
 import pyarrow
 import pyarrow.dataset
 import pyarrow.parquet
 import pytz
-from pydantic.typing import Literal
 
 from feast.data_source import DataSource
 from feast.errors import (

sdk/python/feast/infra/offline_stores/redshift.py

@@ -14,13 +14,13 @@
     Union,
 )
 
+from typing_extensions import Literal
 import numpy as np
 import pandas as pd
 import pyarrow
 import pyarrow as pa
 from dateutil import parser
 from pydantic import StrictStr, root_validator
-from pydantic.typing import Literal
 from pytz import utc
 
 from feast import OnDemandFeatureView, RedshiftSource

sdk/python/feast/infra/offline_stores/snowflake.py

@@ -19,11 +19,11 @@
     cast,
 )
 
+from typing_extensions import Literal
 import numpy as np
 import pandas as pd
 import pyarrow
-from pydantic import Field, StrictStr
-from pydantic.typing import Literal
+from pydantic import ConfigDict, Field, StrictStr
 from pytz import utc
 
 from feast import OnDemandFeatureView
@@ -107,9 +107,7 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel):
 
     convert_timestamp_columns: Optional[bool] = None
     """ Convert timestamp columns on export to a Parquet-supported format """
-
-    class Config:
-        allow_population_by_field_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class SnowflakeOfflineStore(OfflineStore):

sdk/python/feast/infra/online_stores/bigtable.py

@@ -3,12 +3,12 @@
 from concurrent import futures
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple
+from typing_extensions import Literal
 
 import google
 from google.cloud import bigtable
 from google.cloud.bigtable import row_filters
 from pydantic import StrictStr
-from pydantic.typing import Literal
 
 from feast import Entity, FeatureView, utils
 from feast.feature_view import DUMMY_ENTITY_NAME

sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py

@@ -20,7 +20,7 @@
 
 import logging
 from datetime import datetime
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Literal
 
 from cassandra.auth import PlainTextAuthProvider
 from cassandra.cluster import (
@@ -34,7 +34,6 @@
 from cassandra.policies import DCAwareRoundRobinPolicy, TokenAwarePolicy
 from cassandra.query import PreparedStatement
 from pydantic import StrictFloat, StrictInt, StrictStr
-from pydantic.typing import Literal
 
 from feast import Entity, FeatureView, RepoConfig
 from feast.infra.key_encoding_utils import serialize_entity_key

sdk/python/feast/infra/online_stores/contrib/hbase_online_store/hbase.py

@@ -2,9 +2,9 @@
 import struct
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing_extensions import Literal
 
 from happybase import Connection
-from pydantic.typing import Literal
 
 from feast import Entity
 from feast.feature_view import FeatureView

sdk/python/feast/infra/online_stores/datastore.py

@@ -18,9 +18,9 @@
 from queue import Empty, Queue
 from threading import Lock, Thread
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
+from typing_extensions import Literal
 
 from pydantic import PositiveInt, StrictStr
-from pydantic.typing import Literal
 
 from feast import Entity, utils
 from feast.errors import FeastProviderLoginError

sdk/python/feast/infra/online_stores/dynamodb.py

@@ -15,9 +15,9 @@
 import logging
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing_extensions import Literal
 
 from pydantic import StrictBool, StrictStr
-from pydantic.typing import Literal, Union
 
 from feast import Entity, FeatureView, utils
 from feast.infra.infra_object import DYNAMODB_INFRA_OBJECT_CLASS_TYPE, InfraObject

sdk/python/feast/infra/online_stores/redis.py

@@ -27,10 +27,10 @@
     Union,
 )
 
+from typing_extensions import Literal
 import pytz
 from google.protobuf.timestamp_pb2 import Timestamp
 from pydantic import StrictStr
-from pydantic.typing import Literal
 
 from feast import Entity, FeatureView, RepoConfig, utils
 from feast.infra.online_stores.helpers import _mmh3, _redis_key, _redis_key_prefix

sdk/python/feast/infra/online_stores/snowflake.py

@@ -3,10 +3,10 @@
 from binascii import hexlify
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing_extensions import Literal
 
 import pandas as pd
-from pydantic import Field, StrictStr
-from pydantic.schema import Literal
+from pydantic import ConfigDict, Field, StrictStr
 
 from feast.entity import Entity
 from feast.feature_view import FeatureView
@@ -57,9 +57,7 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel):
 
     schema_: Optional[str] = Field("PUBLIC", alias="schema")
     """ Snowflake schema name """
-
-    class Config:
-        allow_population_by_field_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class SnowflakeOnlineStore(OnlineStore):

sdk/python/feast/infra/passthrough_provider.py

@@ -70,7 +70,7 @@ def batch_engine(self) -> BatchMaterializationEngine:
         if self._batch_engine:
             return self._batch_engine
         else:
-            engine_config = self.repo_config._batch_engine_config
+            engine_config = self.repo_config.batch_engine_config
             config_is_dict = False
             if isinstance(engine_config, str):
                 engine_config_type = engine_config

sdk/python/feast/infra/registry/snowflake.py

@@ -7,7 +7,7 @@
 from threading import Lock
 from typing import Any, Callable, List, Optional, Set, Union
 
-from pydantic import Field, StrictStr
+from pydantic import ConfigDict, Field, StrictStr
 from pydantic.schema import Literal
 
 import feast
@@ -103,9 +103,7 @@ class SnowflakeRegistryConfig(RegistryConfig):
 
     schema_: Optional[str] = Field("PUBLIC", alias="schema")
     """ Snowflake schema name """
-
-    class Config:
-        allow_population_by_field_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class SnowflakeRegistry(BaseRegistry):