Merge branch 'main' into migrate_to_hyper_bump

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## Changed
+- Upgrade to Pydantic V2
+
 ## [1.0.2] - 2023-10-10
 
 ### Fixed

gbq/dto.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field, root_validator, validator
+from pydantic import BaseModel, Field, model_validator
 
 
 class StructureType(Enum):
@@ -42,8 +42,14 @@ class BigQueryDataType(Enum):
 
 class TimeDefinition(BaseModel):
     type: TimeType
-    expirationMs: Optional[str]
-    field: Optional[str]
+    expirationMs: Optional[str] = None
+    field: Optional[str] = None
+
+    @model_validator(mode="before")  # type: ignore[arg-type]
+    def str_or_list_(self):
+        if isinstance(self.get("type"), str):
+            self["type"] = TimeType[self.get("type", "").upper()]
+        return self
 
 
 class RangeFieldDefinition(BaseModel):
@@ -66,9 +72,11 @@ class Argument(BaseModel):
     name: str
     data_type: BigQueryDataType
 
-    @validator("data_type", pre=True)
-    def str_or_list_(cls, v):
-        return v.upper()
+    @model_validator(mode="before")  # type: ignore[arg-type]
+    def str_or_list_(self):
+        if isinstance(self.get("data_type"), str):
+            self["data_type"] = BigQueryDataType[self.get("data_type", "").upper()]
+        return self
 
 
 class Structure(BaseModel):
@@ -82,19 +90,21 @@ class Structure(BaseModel):
     type: Union[StructureType, None] = None
     arguments: Union[List[Argument], None] = None
 
-    @root_validator
-    def validate_type(cls, values):
-        if not values.get("type", None):
-            if values["view_query"]:
-                values["type"] = StructureType.view
-            elif values["body"]:
-                values["type"] = StructureType.stored_procedure
+    @model_validator(mode="before")  # type: ignore[arg-type]
+    def validate_type(self):
+        if not self.get("type", None):
+            if self.get("view_query"):
+                self["type"] = StructureType.view
+            elif self.get("body"):
+                self["type"] = StructureType.stored_procedure
             else:
-                values["type"] = StructureType.table
-        return values
-
-    @validator("view_query", "body", pre=True)
-    def str_or_list_(cls, v):
-        if isinstance(v, list) and not [s for s in v if not isinstance(s, str)]:
-            v = "\n".join(v)
-        return v
+                self["type"] = StructureType.table
+        return self
+
+    @model_validator(mode="before")  # type: ignore[arg-type]
+    def str_or_list_body(self):
+        if isinstance(self.get("body"), list) and not [
+            s for s in self.get("body") if not isinstance(s, str)
+        ]:
+            self["body"] = "\n".join(self["body"])
+        return self

requirements-test.txt

@@ -1,6 +1,6 @@
 autoflake==2.2.1
 bandit==1.7.5
-black==23.10.0
+black==23.10.1
 flake8==6.1.0
 isort==5.12.0
 hyper-bump-it==0.5.2; python_version >= '3.9'

requirements.lock

@@ -1,7 +1,8 @@
 # THIS IS AN AUTOGENERATED LOCKFILE. DO NOT EDIT MANUALLY.
+annotated-types==0.6.0
 cachetools==5.3.1
 certifi==2023.7.22
-charset-normalizer==3.3.0
+charset-normalizer==3.3.1
 click==8.1.7
 dataclasses==0.6
 google-api-core==2.12.0
@@ -20,7 +21,8 @@ proto-plus==1.22.3
 protobuf==4.24.4
 pyasn1==0.5.0
 pyasn1-modules==0.3.0
-pydantic==1.10.13
+pydantic==2.4.2
+pydantic_core==2.10.1
 python-dateutil==2.8.2
 requests==2.31.0
 rsa==4.9

requirements.txt

@@ -1,6 +1,6 @@
 dataclasses==0.6
 click==8.1.7
-pydantic<2.0
+pydantic==2.4.2
 
 #bigquery
 google-api-core==2.12.0

setup.cfg

@@ -25,7 +25,7 @@ python_requires = >=3.8
 install_requires =
     typing_extensions~=4.3; python_version < '3.8'
     google-cloud-bigquery
-    pydantic<2.0
+    pydantic>=2.4.0,<3
 packages = find:
 
 [options.package_data]

tests/fixtures.py

@@ -96,7 +96,7 @@ def nested_json_schema_with_time_partition(nested_json_schema):
 def nested_json_schema_with_incorrect_partition(nested_json_schema):
     return {
         "schema": nested_json_schema,
-        "partition": {"type": "time", "definition": {"type": "day"}},
+        "partition": {"type": "time", "definition": {"type": "adasd"}},
         "type": "table",
     }
 

tests/test_bigquery.py

@@ -3,14 +3,16 @@
 from google.cloud import bigquery
 from google.cloud.bigquery import PartitionRange, QueryJob
 from google.cloud.bigquery.routine import RoutineArgument
-from pydantic import ValidationError
 
 from gbq.bigquery import BigQuery
 from gbq.dto import (
+    Argument,
+    BigQueryDataType,
     Partition,
     RangeDefinition,
     RangeFieldDefinition,
     Structure,
+    StructureType,
     TimeDefinition,
 )
 from gbq.exceptions import GbqException
@@ -132,7 +134,7 @@ def test_create_or_update_structure_with_incorrect_partition(
     bq, nested_json_schema_with_incorrect_partition, table
 ):
     bq.bq_client.get_table.side_effect = NotFound("")
-    with pytest.raises(ValidationError):
+    with pytest.raises(KeyError):
         bq.create_or_update_structure(
             "project",
             "dataset",
@@ -487,3 +489,22 @@ def test_execute_raise_no_exception(bq):
     response = bq.execute(query=query)
 
     assert isinstance(response, QueryJob)
+
+
+def test_argument_class_data_type():
+    input_json = {"name": "test", "data_type": "string"}
+    expected = Argument(**input_json)
+    assert expected.name == "test"
+    assert expected.data_type == BigQueryDataType.STRING
+
+
+def test_structure_validate_type():
+    input_json = {"view_query": "test"}
+    expected = Structure(**input_json)
+    assert expected.type == StructureType.view
+
+
+def test_structure_str_or_list_body():
+    input_json = {"body": ["test", "test1", "test2"]}
+    expected = Structure(**input_json)
+    assert expected.body == "test\ntest1\ntest2"