Skip to content

Commit

Permalink
Fix up PR for Pydantic 2.0
Browse files Browse the repository at this point in the history
  • Loading branch information
jmchilton committed Jan 15, 2024
1 parent 9e62358 commit 7d98cb8
Show file tree
Hide file tree
Showing 5 changed files with 39 additions and 44 deletions.
6 changes: 4 additions & 2 deletions lib/galaxy/objectstore/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1194,7 +1194,7 @@ def to_dict(self, object_store_uris: Optional[Set[str]] = None) -> Dict[str, Any
object_store_config = self.user_object_store_resolver.resolve_object_store_uri_config(
user_object_store_uri
)
dynamic_object_store_as_dict = object_store_config.dict()
dynamic_object_store_as_dict = object_store_config.model_dump()
dynamic_object_store_as_dict["id"] = user_object_store_uri
dynamic_object_store_as_dict["weight"] = 0
# these are all forward facing object stores...
Expand Down Expand Up @@ -1622,7 +1622,9 @@ class GalaxyConfigAdapter:
)
assert issubclass(objectstore_class, ConcreteObjectStore)
return objectstore_class(
config=GalaxyConfigAdapter(), config_dict=object_store_configuration.dict(), **objectstore_constructor_kwds
config=GalaxyConfigAdapter(),
config_dict=object_store_configuration.model_dump(),
**objectstore_constructor_kwds,
)


Expand Down
12 changes: 6 additions & 6 deletions lib/galaxy/objectstore/templates/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,26 +64,26 @@ def from_app_config(config: AppConfigProtocol, vault_configured=None) -> "Config

@property
def configuration_uses_secrets(self) -> bool:
templates = self.catalog.__root__
templates = self.catalog.root
for template in templates:
if template.secrets and len(template.secrets) > 0:
return True
return False

@property
def summaries(self) -> ObjectStoreTemplateSummaries:
templates = self.catalog.__root__
templates = self.catalog.root
summaries = []
for template in templates:
template_dict = template.dict()
template_dict = template.model_dump()
configuration = template_dict.pop("configuration")
stored_badges = configuration["badges"] or []
object_store_type = configuration["type"]
badges = serialize_badges(stored_badges, False, True, True, object_store_type in ["azure", "s3"])
template_dict["badges"] = badges
template_dict["type"] = object_store_type
summaries.append(template_dict)
return ObjectStoreTemplateSummaries.parse_obj(summaries)
return ObjectStoreTemplateSummaries.model_validate(summaries)

def find_template(self, instance_reference: TemplateReference) -> ObjectStoreTemplate:
"""Find the corresponding template and throw ObjectNotFound if not available."""
Expand All @@ -92,7 +92,7 @@ def find_template(self, instance_reference: TemplateReference) -> ObjectStoreTem
return self.find_template_by(template_id, template_version)

def find_template_by(self, template_id: str, template_version: int) -> ObjectStoreTemplate:
templates = self.catalog.__root__
templates = self.catalog.root

for template in templates:
if template.id == template_id and template.version == template_version:
Expand All @@ -119,7 +119,7 @@ def validate(self, instance: InstanceDefinition):

def raw_config_to_catalog(raw_config: List[RawTemplateConfig]) -> ObjectStoreTemplateCatalog:
effective_root = _apply_syntactic_sugar(raw_config)
return ObjectStoreTemplateCatalog.parse_obj(effective_root)
return ObjectStoreTemplateCatalog.model_validate(effective_root)


def _apply_syntactic_sugar(raw_templates: List[RawTemplateConfig]) -> List[RawTemplateConfig]:
Expand Down
41 changes: 17 additions & 24 deletions lib/galaxy/objectstore/templates/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@
from jinja2.nativetypes import NativeEnvironment
from pydantic import (
BaseModel,
Extra,
ConfigDict,
RootModel,
)
from typing_extensions import Literal

Expand All @@ -22,8 +23,7 @@


class StrictModel(BaseModel):
class Config:
extra = Extra.forbid
model_config = ConfigDict(extra="forbid")


ObjectStoreTemplateVariableType = Literal["string", "boolean", "integer"]
Expand Down Expand Up @@ -59,14 +59,14 @@ class S3ObjectStoreTemplateConfiguration(StrictModel):
type: Literal["s3"]
auth: S3AuthTemplate
bucket: S3BucketTemplate
badges: BadgeList
badges: BadgeList = None


class S3ObjectStoreConfiguration(StrictModel):
type: Literal["s3"]
auth: S3Auth
bucket: S3Bucket
badges: BadgeList
badges: BadgeList = None


class AzureAuthTemplate(StrictModel):
Expand All @@ -91,26 +91,26 @@ class AzureObjectStoreTemplateConfiguration(StrictModel):
type: Literal["azure_blob"]
auth: AzureAuthTemplate
container: AzureContainerTemplate
badges: BadgeList
badges: BadgeList = None


class AzureObjectStoreConfiguration(StrictModel):
type: Literal["azure_blob"]
auth: AzureAuth
container: AzureContainer
badges: BadgeList
badges: BadgeList = None


class DiskObjectStoreTemplateConfiguration(StrictModel):
type: Literal["disk"]
files_dir: Union[str, TemplateExpansion]
badges: BadgeList
badges: BadgeList = None


class DiskObjectStoreConfiguration(StrictModel):
type: Literal["disk"]
files_dir: str
badges: BadgeList
badges: BadgeList = None


class S3ConnectionTemplate(StrictModel):
Expand All @@ -132,15 +132,15 @@ class GenericS3ObjectStoreTemplateConfiguration(StrictModel):
auth: S3AuthTemplate
bucket: S3BucketTemplate
connection: S3ConnectionTemplate
badges: BadgeList
badges: BadgeList = None


class GenericS3ObjectStoreConfiguration(StrictModel):
type: Literal["generic_s3"]
auth: S3Auth
bucket: S3Bucket
connection: S3Connection
badges: BadgeList
badges: BadgeList = None


ObjectStoreTemplateConfiguration = Union[
Expand Down Expand Up @@ -188,29 +188,22 @@ class ObjectStoreTemplateBase(StrictModel):
# template by hiding but keep it in the catalog for backward
# compatibility for users with existing stores of that template.
hidden: bool = False
variables: Optional[List[ObjectStoreTemplateVariable]]
secrets: Optional[List[ObjectStoreTemplateSecret]]
variables: Optional[List[ObjectStoreTemplateVariable]] = None
secrets: Optional[List[ObjectStoreTemplateSecret]] = None


class ObjectStoreTemplateSummary(ObjectStoreTemplateBase):
badges: List[BadgeDict]
badges: List[BadgeDict] = None
type: ObjectStoreTemplateType


class ObjectStoreTemplate(ObjectStoreTemplateBase):
configuration: ObjectStoreTemplateConfiguration


class ObjectStoreTemplateCatalog(StrictModel):
"""Represents a collection of ObjectStoreTemplates."""
ObjectStoreTemplateCatalog = RootModel[List[ObjectStoreTemplate]]

__root__: List[ObjectStoreTemplate]


class ObjectStoreTemplateSummaries(StrictModel):
"""Represents a collection of ObjectStoreTemplate summaries."""

__root__: List[ObjectStoreTemplateSummary]
ObjectStoreTemplateSummaries = RootModel[List[ObjectStoreTemplateSummary]]


def template_to_configuration(
Expand All @@ -233,7 +226,7 @@ def expand_template(_, key, value):
return key, template.render(**template_variables)
return key, value

raw_config = remap(configuration_template.dict(), visit=expand_template)
raw_config = remap(configuration_template.model_dump(), visit=expand_template)
return to_configuration_object(raw_config)


Expand Down
2 changes: 1 addition & 1 deletion test/unit/objectstore/test_template_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ def test_manager(tmpdir):
templates = ConfiguredObjectStoreTemplates.from_app_config(config)
summaries = templates.summaries
assert summaries
assert len(summaries.__root__) == 2
assert len(summaries.root) == 2


def test_manager_throws_exception_if_vault_is_required_but_configured(tmpdir):
Expand Down
22 changes: 11 additions & 11 deletions test/unit/objectstore/test_template_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,8 @@

def test_parsing_simple_s3():
template_library = _parse_template_library(LIBRARY_1)
assert len(template_library.__root__) == 1
s3_template = template_library.__root__[0]
assert len(template_library.root) == 1
s3_template = template_library.root[0]
assert s3_template.description == "An Amazon S3 Bucket"
configuration_obj = template_to_configuration(
s3_template,
Expand All @@ -57,7 +57,7 @@ def test_parsing_simple_s3():

# expanded configuration should validate with template expansions...
assert isinstance(configuration_obj, S3ObjectStoreConfiguration)
configuration = configuration_obj.dict()
configuration = configuration_obj.model_dump()

assert configuration["type"] == "s3"
assert configuration["auth"]["access_key"] == "sec1"
Expand Down Expand Up @@ -102,8 +102,8 @@ def test_parsing_simple_s3():

def test_parsing_generic_s3():
template_library = _parse_template_library(LIBRARY_GENERIC_S3)
assert len(template_library.__root__) == 1
s3_template = template_library.__root__[0]
assert len(template_library.root) == 1
s3_template = template_library.root[0]
assert s3_template.description == "A MinIO bucket connected using a generic S3 object store."
configuration_obj = template_to_configuration(
s3_template,
Expand All @@ -117,7 +117,7 @@ def test_parsing_generic_s3():

# expanded configuration should validate with template expansions...
assert isinstance(configuration_obj, GenericS3ObjectStoreConfiguration)
configuration = configuration_obj.dict()
configuration = configuration_obj.model_dump()

assert configuration["type"] == "generic_s3"
assert configuration["auth"]["access_key"] == "sec1"
Expand Down Expand Up @@ -149,9 +149,9 @@ def test_parsing_generic_s3():

def test_parsing_multiple_posix():
template_library = _parse_template_library(LIBRARY_2)
assert len(template_library.__root__) == 2
general_template = template_library.__root__[0]
secure_template = template_library.__root__[1]
assert len(template_library.root) == 2
general_template = template_library.root[0]
secure_template = template_library.root[1]

assert general_template.version == 0
assert secure_template.version == 0
Expand Down Expand Up @@ -191,8 +191,8 @@ def test_parsing_multiple_posix():

def test_parsing_azure():
template_library = _parse_template_library(LIBRARY_AZURE_CONTAINER)
assert len(template_library.__root__) == 1
azure_template = template_library.__root__[0]
assert len(template_library.root) == 1
azure_template = template_library.root[0]
assert azure_template.description == "An Azure Container"
configuration_obj = template_to_configuration(
azure_template,
Expand Down

0 comments on commit 7d98cb8

Please sign in to comment.