import re from enum import Enum from typing import Any, Dict, Iterable, Mapping, Optional, Set from urllib.parse import urlparse from wandb_gql import gql from wandb_graphql.language import ast, visitor from wandb._iterutils import one from wandb.sdk.artifacts._validators import is_artifact_registry_project from wandb.sdk.internal.internal_api import Api as InternalApi def parse_s3_url_to_s3_uri(url) -> str: """Convert an S3 HTTP(S) URL to an S3 URI. Arguments: url (str): The S3 URL to convert, in the format 'http(s)://.s3..amazonaws.com/'. or 'http(s)://.s3.amazonaws.com/' Returns: str: The corresponding S3 URI in the format 's3:///'. Raises: ValueError: If the provided URL is not a valid S3 URL. """ # Regular expression to match S3 URL pattern s3_pattern = r"^https?://.*s3.*amazonaws\.com.*" parsed_url = urlparse(url) # Check if it's an S3 URL match = re.match(s3_pattern, parsed_url.geturl()) if not match: raise ValueError("Invalid S3 URL") # Extract bucket name and key bucket_name, *_ = parsed_url.netloc.split(".") key = parsed_url.path.lstrip("/") # Construct the S3 URI s3_uri = f"s3://{bucket_name}/{key}" return s3_uri class PathType(Enum): """We have lots of different paths users pass in to fetch artifacts, projects, etc. This enum is used for specifying what format the path is in given a string path. """ PROJECT = "PROJECT" ARTIFACT = "ARTIFACT" def parse_org_from_registry_path(path: str, path_type: PathType) -> str: """Parse the org from a registry path. Essentially fetching the "entity" from the path but for Registries the entity is actually the org. Args: path (str): The path to parse. Can be a project path / or or an artifact path like // or / or path_type (PathType): The type of path to parse. """ parts = path.split("/") expected_parts = 3 if path_type == PathType.ARTIFACT else 2 if len(parts) >= expected_parts: org, project = parts[:2] if is_artifact_registry_project(project): return org return "" def fetch_org_from_settings_or_entity( settings: dict, default_entity: Optional[str] = None ) -> str: """Fetch the org from either the settings or deriving it from the entity. Returns the org from the settings if available. If no org is passed in or set, the entity is used to fetch the org. Args: organization (str | None): The organization to fetch the org for. settings (dict): The settings to fetch the org for. default_entity (str | None): The default entity to fetch the org for. """ if (organization := settings.get("organization")) is None: # Fetch the org via the Entity. Won't work if default entity is a personal entity and belongs to multiple orgs entity = settings.get("entity") or default_entity if entity is None: raise ValueError( "No entity specified and can't fetch organization from the entity" ) entity_orgs = InternalApi()._fetch_orgs_and_org_entities_from_entity(entity) entity_org = one( entity_orgs, too_short=ValueError( "No organizations found for entity. Please specify an organization in the settings." ), too_long=ValueError( "Multiple organizations found for entity. Please specify an organization in the settings." ), ) organization = entity_org.display_name return organization class _GQLCompatRewriter(visitor.Visitor): """GraphQL AST visitor to rewrite queries/mutations to be compatible with older server versions.""" omit_variables: Set[str] omit_fragments: Set[str] omit_fields: Set[str] rename_fields: Dict[str, str] def __init__( self, omit_variables: Optional[Iterable[str]] = None, omit_fragments: Optional[Iterable[str]] = None, omit_fields: Optional[Iterable[str]] = None, rename_fields: Optional[Mapping[str, str]] = None, ): self.omit_variables = set(omit_variables or ()) self.omit_fragments = set(omit_fragments or ()) self.omit_fields = set(omit_fields or ()) self.rename_fields = dict(rename_fields or {}) def enter_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any: # noqa: N802 if node.variable.name.value in self.omit_variables: return visitor.REMOVE # return node def enter_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any: # noqa: N802 # For context, note that e.g.: # # {description: $description # ...} # # Is parsed as: # # ObjectValue(fields=[ # ObjectField(name=Name(value='description'), value=Variable(name=Name(value='description'))), # ...]) if ( isinstance(var := node.value, ast.Variable) and var.name.value in self.omit_variables ): return visitor.REMOVE def enter_Argument(self, node: ast.Argument, *_, **__) -> Any: # noqa: N802 if node.name.value in self.omit_variables: return visitor.REMOVE def enter_FragmentDefinition(self, node: ast.FragmentDefinition, *_, **__) -> Any: # noqa: N802 if node.name.value in self.omit_fragments: return visitor.REMOVE def enter_FragmentSpread(self, node: ast.FragmentSpread, *_, **__) -> Any: # noqa: N802 if node.name.value in self.omit_fragments: return visitor.REMOVE def enter_Field(self, node: ast.Field, *_, **__) -> Any: # noqa: N802 if node.name.value in self.omit_fields: return visitor.REMOVE if new_name := self.rename_fields.get(node.name.value): node.name.value = new_name return node def leave_Field(self, node: ast.Field, *_, **__) -> Any: # noqa: N802 # If the field had a selection set, but now it's empty, remove the field entirely if (node.selection_set is not None) and (not node.selection_set.selections): return visitor.REMOVE def gql_compat( request_string: str, omit_variables: Optional[Iterable[str]] = None, omit_fragments: Optional[Iterable[str]] = None, omit_fields: Optional[Iterable[str]] = None, rename_fields: Optional[Mapping[str, str]] = None, ) -> ast.Document: """Rewrite a GraphQL request string to ensure compatibility with older server versions. Args: request_string (str): The GraphQL request string to rewrite. omit_variables (Iterable[str] | None): Names of variables to remove from the request string. omit_fragments (Iterable[str] | None): Names of fragments to remove from the request string. omit_fields (Iterable[str] | None): Names of fields to remove from the request string. rename_fields (Mapping[str, str] | None): A mapping of fields to rename in the request string, given as `{old_name -> new_name}`. Returns: str: Modified GraphQL request string with fragments on omitted types removed. """ # Parse the request into a GraphQL AST doc = gql(request_string) if not (omit_variables or omit_fragments or omit_fields or rename_fields): return doc # Visit the AST with our visitor to filter out unwanted fragments rewriter = _GQLCompatRewriter( omit_variables=omit_variables, omit_fragments=omit_fragments, omit_fields=omit_fields, rename_fields=rename_fields, ) return visitor.visit(doc, rewriter)