build: Relax pydantic version requirement

setup.cfg

@@ -32,7 +32,7 @@ install_requires =
     cryptography>=3.4.0
     httpcore<0.17.3
     httpx[http2]==0.24.0
-    pydantic[dotenv]>=1.8.2,<2.0
+    pydantic[dotenv]>=1.8.2,<3.0.0
     python-dateutil>=2.8.2
     readerwriterlock>=1.0.9
     sqlparse>=0.4.2
@@ -51,7 +51,7 @@ ciso8601 =
 dev =
     allure-pytest==2.*
     devtools==0.7.0
-    mypy==0.910
+    mypy==1.*
     pre-commit==2.15.0
     pyfakefs>=4.5.3
     pytest==7.2.0
@@ -71,6 +71,7 @@ firebolt = py.typed
 disallow_untyped_defs = True
 show_error_codes = True
 files = src/
+plugins = pydantic.mypy
 
 [flake8]
 exclude = tests/*

src/firebolt/async_db/cursor.py

@@ -395,7 +395,7 @@ def __exit__(
         self.close()
 
     def __await__(self) -> Iterator:
-        pass
+        yield None
 
     async def __aexit__(
         self, exc_type: type, exc_val: Exception, exc_tb: TracebackType

src/firebolt/client/auth/base.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from time import time
 from typing import AsyncGenerator, Generator, Optional
 
@@ -44,7 +45,7 @@ def copy(self) -> "Auth":
         Returns:
             Auth: Auth object
         """
-        return Auth(self._use_token_cache)
+        return self.__class__(self._use_token_cache)
 
     @property
     def token(self) -> Optional[str]:
@@ -103,6 +104,7 @@ def _cache_token(self) -> None:
         if self._token and self._expires:
             self._token_storage.cache_token(self._token, self._expires)
 
+    @abstractmethod
     def get_new_token_generator(self) -> Generator[Request, Response, None]:
         """Generate requests needed to create a new token session."""
 

src/firebolt/model/V1/__init__.py

@@ -5,9 +5,12 @@
 
 
 class FireboltBaseModel(BaseModel):
+
+    # Using Pydantic 1.* config class for backwards compatibility
     class Config:
-        allow_population_by_field_name = True
         extra = "forbid"
+        allow_population_by_field_name = True  # Pydantic 1.8
+        populate_by_name = True  # Pydantic 2.0
 
     def jsonable_dict(self, *args: Any, **kwargs: Any) -> dict:
         """

src/firebolt/model/V1/binding.py

@@ -21,13 +21,13 @@ class Binding(FireboltBaseModel):
     is_default_engine: bool = Field(alias="engine_is_default")
 
     # optional
-    current_status: Optional[str]
-    health_status: Optional[str]
-    create_time: Optional[datetime]
-    create_actor: Optional[str]
-    last_update_time: Optional[datetime]
-    last_update_actor: Optional[str]
-    desired_status: Optional[str]
+    current_status: Optional[str] = None
+    health_status: Optional[str] = None
+    create_time: Optional[datetime] = None
+    create_actor: Optional[str] = None
+    last_update_time: Optional[datetime] = None
+    last_update_actor: Optional[str] = None
+    desired_status: Optional[str] = None
 
     @property
     def database_id(self) -> str:

src/firebolt/model/V1/database.py

@@ -42,24 +42,24 @@ class Database(FireboltBaseModel):
     _service: DatabaseService = PrivateAttr()
 
     # required
-    name: str = Field(min_length=1, max_length=255, regex=r"^[0-9a-zA-Z_]+$")
+    name: str = Field(min_length=1, max_length=255, pattern=r"^[0-9a-zA-Z_]+$")
     compute_region_key: RegionKey = Field(alias="compute_region_id")
 
     # optional
-    database_key: Optional[DatabaseKey] = Field(None, alias="id")
-    description: Optional[str] = Field(None, max_length=255)
-    emoji: Optional[str] = Field(None, max_length=255)
-    current_status: Optional[str]
-    health_status: Optional[str]
-    data_size_full: Optional[int]
-    data_size_compressed: Optional[int]
-    is_system_database: Optional[bool]
-    storage_bucket_name: Optional[str]
-    create_time: Optional[datetime]
-    create_actor: Optional[str]
-    last_update_time: Optional[datetime]
-    last_update_actor: Optional[str]
-    desired_status: Optional[str]
+    database_key: Optional[DatabaseKey] = Field(default=None, alias="id")
+    description: Optional[str] = Field(default=None, max_length=255)
+    emoji: Optional[str] = Field(default=None, max_length=255)
+    current_status: Optional[str] = None
+    health_status: Optional[str] = None
+    data_size_full: Optional[int] = None
+    data_size_compressed: Optional[int] = None
+    is_system_database: Optional[bool] = None
+    storage_bucket_name: Optional[str] = None
+    create_time: Optional[datetime] = None
+    create_actor: Optional[str] = None
+    last_update_time: Optional[datetime] = None
+    last_update_actor: Optional[str] = None
+    desired_status: Optional[str] = None
 
     @classmethod
     def parse_obj_with_service(
@@ -142,6 +142,7 @@ class _DatabaseUpdateRequest(FireboltBaseModel):
 
         self.description = description
 
+        assert self.database_id is not None, "Database must have database_id"
         logger.info(
             f"Updating Database (database_id={self.database_id}, "
             f"name={self.name}, description={self.description})"

src/firebolt/model/V1/engine.py

@@ -50,7 +50,7 @@ class EngineSettings(FireboltBaseModel):
     """
 
     preset: str
-    auto_stop_delay_duration: str = Field(regex=r"^[0-9]+[sm]$|^0$")
+    auto_stop_delay_duration: str = Field(pattern=r"^[0-9]+[sm]$|^0$")
     minimum_logging_level: str
     is_read_only: bool
     warm_up: str
@@ -97,30 +97,30 @@ class Engine(FireboltBaseModel):
     _service: EngineService = PrivateAttr()
 
     # required
-    name: str = Field(min_length=1, max_length=255, regex=r"^[0-9a-zA-Z_]+$")
+    name: str = Field(min_length=1, max_length=255, pattern=r"^[0-9a-zA-Z_]+$")
     compute_region_key: RegionKey = Field(alias="compute_region_id")
     settings: EngineSettings
 
     # optional
     key: Optional[EngineKey] = Field(None, alias="id")
-    description: Optional[str]
-    emoji: Optional[str]
-    current_status: Optional[EngineStatus]
-    current_status_summary: Optional[EngineStatusSummary]
+    description: Optional[str] = None
+    emoji: Optional[str] = None
+    current_status: Optional[EngineStatus] = None
+    current_status_summary: Optional[EngineStatusSummary] = None
     latest_revision_key: Optional[EngineRevisionKey] = Field(
         None, alias="latest_revision_id"
     )
-    endpoint: Optional[str]
+    endpoint: Optional[str] = None
     endpoint_serving_revision_key: Optional[EngineRevisionKey] = Field(
         None, alias="endpoint_serving_revision_id"
     )
-    create_time: Optional[datetime]
-    create_actor: Optional[str]
-    last_update_time: Optional[datetime]
-    last_update_actor: Optional[str]
-    last_use_time: Optional[datetime]
-    desired_status: Optional[str]
-    health_status: Optional[str]
+    create_time: Optional[datetime] = None
+    create_actor: Optional[str] = None
+    last_update_time: Optional[datetime] = None
+    last_update_actor: Optional[str] = None
+    last_use_time: Optional[datetime] = None
+    desired_status: Optional[str] = None
+    health_status: Optional[str] = None
     endpoint_desired_revision_key: Optional[EngineRevisionKey] = Field(
         None, alias="endpoint_desired_revision_id"
     )

src/firebolt/model/V1/provider.py

@@ -11,6 +11,6 @@ class Provider(FireboltBaseModel, frozen=True):  # type: ignore
     name: str
 
     # optional
-    create_time: Optional[datetime]
-    display_name: Optional[str]
-    last_update_time: Optional[datetime]
+    create_time: Optional[datetime] = None
+    display_name: Optional[str] = None
+    last_update_time: Optional[datetime] = None

src/firebolt/model/V1/region.py

@@ -16,6 +16,6 @@ class Region(FireboltBaseModel):
     name: str
 
     # optional
-    display_name: Optional[str]
-    create_time: Optional[datetime]
-    last_update_time: Optional[datetime]
+    display_name: Optional[str] = None
+    create_time: Optional[datetime] = None
+    last_update_time: Optional[datetime] = None

src/firebolt/service/V1/binding.py

@@ -125,6 +125,7 @@ def create(
             f"to Database (database_id={database.database_id}, "
             f"name={database.name})"
         )
+        assert database.database_id is not None, "Database must have database_id"
         binding = Binding(
             binding_key=BindingKey(
                 account_id=self.account_id,

src/firebolt/service/V1/region.py

@@ -61,6 +61,13 @@ def get_by_key(self, key: RegionKey) -> Region:
     def get_by_id(self, id_: str) -> Region:
         """Get an AWS region by region_id."""
 
+        # error if provider_id is not set
+        if not self.resource_manager.provider_id:
+            raise ValueError(
+                "provider_id parameter must be set when initializing "
+                "the resource manager."
+            )
+
         return self.get_by_key(
             RegionKey(provider_id=self.resource_manager.provider_id, region_id=id_)
         )