feat: Update the pydantic version to v2.

Signed-off-by: Shuchu Han <shuchu.han@gmail.com>
feast-dev · shuchu · Feb 15, 2024 · Feb 7, 2024 · Feb 8, 2024 · Feb 8, 2024
commit 6a19c505073cf0424e025e9d5aff2f035d76feef
@@ -1,5 +1,6 @@
+from typing import Literal
+
 from pydantic import StrictBool, StrictStr
-from pydantic.typing import Literal
 
 from feast.infra.feature_servers.base_config import BaseFeatureServerConfig
 

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

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

@@ -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

@@ -10,6 +10,7 @@
     Dict,
     Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     Union,
@@ -19,8 +20,7 @@
 import pandas as pd
 import pyarrow
 import pyarrow.parquet
-from pydantic import ConstrainedStr, StrictStr, validator
-from pydantic.typing import Literal
+from pydantic import StrictStr, field_validator
 from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
 
 from feast import flags_helper
@@ -72,13 +72,6 @@ def get_http_client_info():
     return http_client_info.ClientInfo(user_agent=get_user_agent())
 
 
-class BigQueryTableCreateDisposition(ConstrainedStr):
-    """Custom constraint for table_create_disposition. To understand more, see:
-    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationLoad.FIELDS.create_disposition"""
-
-    values = {"CREATE_NEVER", "CREATE_IF_NEEDED"}
-
-
 class BigQueryOfflineStoreConfig(FeastConfigBaseModel):
     """Offline store config for GCP BigQuery"""
 
@@ -102,10 +95,15 @@ class BigQueryOfflineStoreConfig(FeastConfigBaseModel):
     gcs_staging_location: Optional[str] = None
     """ (optional) GCS location used for offloading BigQuery results as parquet files."""
 
-    table_create_disposition: Optional[BigQueryTableCreateDisposition] = None
-    """ (optional) Specifies whether the job is allowed to create new tables. The default value is CREATE_IF_NEEDED."""
+    table_create_disposition: Literal[
+        "CREATE_NEVER", "CREATE_IF_NEEDED"
+    ] = "CREATE_IF_NEEDED"
+    """ (optional) Specifies whether the job is allowed to create new tables. The default value is CREATE_IF_NEEDED.
+    Custom constraint for table_create_disposition. To understand more, see:
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationLoad.FIELDS.create_disposition
+    """
 
-    @validator("billing_project_id")
+    @field_validator("billing_project_id")
     def project_id_exists(cls, v, values, **kwargs):
         if v and not values["project_id"]:
             raise ValueError(

@@ -8,6 +8,7 @@
     Dict,
     Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     Union,
@@ -18,7 +19,6 @@
 import pyarrow
 import pyarrow as pa
 from pydantic import StrictStr
-from pydantic.typing import Literal
 from pytz import utc
 
 from feast import OnDemandFeatureView

@@ -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, Literal, Optional, Set, Tuple, Union
 
 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

@@ -9,6 +9,7 @@
     Iterator,
     KeysView,
     List,
+    Literal,
     Optional,
     Tuple,
     Union,
@@ -19,7 +20,6 @@
 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

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 import pyarrow
-from pydantic import Field, FilePath, SecretStr, StrictBool, StrictStr, root_validator
+from pydantic import Field, FilePath, SecretStr, StrictBool, StrictStr, model_validator
 from trino.auth import (
     BasicAuthentication,
     CertificateAuthentication,
@@ -98,14 +98,14 @@ class AuthConfig(FeastConfigBaseModel):
     type: Literal["kerberos", "basic", "jwt", "oauth2", "certificate"]
     config: Optional[Dict[StrictStr, Any]]
 
-    @root_validator
-    def config_only_nullable_for_oauth2(cls, values):
-        auth_type = values["type"]
-        auth_config = values["config"]
+    @model_validator(mode="after")
+    def config_only_nullable_for_oauth2(self):
+        auth_type = self.type
+        auth_config = self.config
         if auth_type != "oauth2" and auth_config is None:
             raise ValueError(f"config cannot be null for auth type '{auth_type}'")
 
-        return values
+        return self
 
     def to_trino_auth(self):
         auth_type = self.type

@@ -2,15 +2,14 @@
 import uuid
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
 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 (

@@ -9,6 +9,7 @@
     Dict,
     Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     Union,
@@ -19,8 +20,7 @@
 import pyarrow
 import pyarrow as pa
 from dateutil import parser
-from pydantic import StrictStr, root_validator
-from pydantic.typing import Literal
+from pydantic import StrictStr, model_validator
 from pytz import utc
 
 from feast import OnDemandFeatureView, RedshiftSource
@@ -72,16 +72,16 @@ class RedshiftOfflineStoreConfig(FeastConfigBaseModel):
     iam_role: StrictStr
     """ IAM Role for Redshift, granting it access to S3 """
 
-    @root_validator
-    def require_cluster_and_user_or_workgroup(cls, values):
+    @model_validator(mode="after")
+    def require_cluster_and_user_or_workgroup(self):
         """
         Provisioned Redshift clusters:  Require cluster_id and user, ignore workgroup
         Serverless Redshift:  Require workgroup, ignore cluster_id and user
         """
         cluster_id, user, workgroup = (
-            values.get("cluster_id"),
-            values.get("user"),
-            values.get("workgroup"),
+            self.cluster_id,
+            self.user,
+            self.workgroup,
         )
         if not (cluster_id and user) and not workgroup:
             raise ValueError(
@@ -90,7 +90,7 @@ def require_cluster_and_user_or_workgroup(cls, values):
         elif cluster_id and workgroup:
             raise ValueError("cannot specify both cluster_id and workgroup")
 
-        return values
+        return self
 
 
 class RedshiftOfflineStore(OfflineStore):

@@ -14,6 +14,7 @@
     Dict,
     Iterator,
     List,
+    Literal,
     Optional,
     Tuple,
     Union,
@@ -23,8 +24,7 @@
 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
@@ -119,9 +119,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):

@@ -2,13 +2,12 @@
 import logging
 from concurrent import futures
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Set, Tuple
 
 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

@@ -20,7 +20,17 @@
 
 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,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+)
 
 from cassandra.auth import PlainTextAuthProvider
 from cassandra.cluster import (
@@ -34,7 +44,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

@@ -1,12 +1,11 @@
 import calendar
 import struct
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple
 
 from happybase import ConnectionPool
 from happybase.connection import DEFAULT_PROTOCOL, DEFAULT_TRANSPORT
 from pydantic import StrictStr
-from pydantic.typing import Literal
 
 from feast import Entity
 from feast.feature_view import FeatureView

@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple
 
 import pymysql
 import pytz
@@ -23,7 +23,7 @@ class MySQLOnlineStoreConfig(FeastConfigBaseModel):
     NOTE: The class *must* end with the `OnlineStoreConfig` suffix.
     """
 
-    type = "mysql"
+    type: Literal["mysql"] = "mysql"
 
     host: Optional[StrictStr] = None
     user: Optional[StrictStr] = None

@@ -2,14 +2,13 @@
 import logging
 from collections import defaultdict
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple
 
 import psycopg2
 import pytz
 from psycopg2 import sql
 from psycopg2.extras import execute_values
 from psycopg2.pool import SimpleConnectionPool
-from pydantic.schema import Literal
 
 from feast import Entity
 from feast.feature_view import FeatureView

@@ -17,10 +17,19 @@
 from multiprocessing.pool import ThreadPool
 from queue import Empty, Queue
 from threading import Lock, Thread
-from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+)
 
 from pydantic import PositiveInt, StrictStr
-from pydantic.typing import Literal
 
 from feast import Entity, utils
 from feast.errors import FeastProviderLoginError

@@ -14,10 +14,9 @@
 import itertools
 import logging
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
 
 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