File size: 7,893 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
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)://<bucket>.s3.<region>.amazonaws.com/<key>'.
or 'http(s)://<bucket>.s3.amazonaws.com/<key>'
Returns:
str: The corresponding S3 URI in the format 's3://<bucket>/<key>'.
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 <entity>/<project> or <project> or an
artifact path like <entity>/<project>/<artifact> or <project>/<artifact> or <artifact>
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)
|