diff --git a/CHANGELOG.md b/CHANGELOG.md index 9eef6ff7..984b30f9 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/gbq/dto.py b/gbq/dto.py index c51b7237..b5235afa 100644 --- a/gbq/dto.py +++ b/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 diff --git a/requirements-test.txt b/requirements-test.txt index edfd802e..f028a62a 100644 --- a/requirements-test.txt +++ b/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' diff --git a/requirements.lock b/requirements.lock index 607ee173..0af89f18 100644 --- a/requirements.lock +++ b/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 diff --git a/requirements.txt b/requirements.txt index 90594928..e8c1bcbe 100644 --- a/requirements.txt +++ b/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 diff --git a/setup.cfg b/setup.cfg index 2cb21b15..6dddd9f4 100644 --- a/setup.cfg +++ b/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] diff --git a/tests/fixtures.py b/tests/fixtures.py index 904500e8..564bc971 100644 --- a/tests/fixtures.py +++ b/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", } diff --git a/tests/test_bigquery.py b/tests/test_bigquery.py index 46a6a0a5..0306ae7c 100644 --- a/tests/test_bigquery.py +++ b/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"