acryl-datahub 1.0.0.2rc4__py3-none-any.whl → 1.0.0.3__py3-none-any.whl

datahub/sdk/mlmodelgroup.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, List, Optional, Sequence, Type, Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mce_builder import DEFAULT_ENV
+from datahub.metadata.schema_classes import (
+    AspectBag,
+    MLModelGroupPropertiesClass,
+)
+from datahub.metadata.urns import DataProcessInstanceUrn, MlModelGroupUrn, Urn
+from datahub.sdk._shared import (
+    DomainInputType,
+    HasDomain,
+    HasInstitutionalMemory,
+    HasOwnership,
+    HasPlatformInstance,
+    HasTags,
+    HasTerms,
+    LinksInputType,
+    OwnersInputType,
+    TagsInputType,
+    TermsInputType,
+    make_time_stamp,
+    parse_time_stamp,
+)
+from datahub.sdk.entity import Entity, ExtraAspectsType
+
+
+class MLModelGroup(
+    HasPlatformInstance,
+    HasOwnership,
+    HasInstitutionalMemory,
+    HasTags,
+    HasTerms,
+    HasDomain,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[MlModelGroupUrn]:
+        return MlModelGroupUrn
+
+    def __init__(
+        self,
+        id: str,
+        platform: str,
+        name: Optional[str] = "",
+        platform_instance: Optional[str] = None,
+        env: str = DEFAULT_ENV,
+        # Model group properties
+        description: Optional[str] = None,
+        display_name: Optional[str] = None,
+        external_url: Optional[str] = None,
+        custom_properties: Optional[Dict[str, str]] = None,
+        created: Optional[datetime] = None,
+        last_modified: Optional[datetime] = None,
+        # Standard aspects
+        owners: Optional[OwnersInputType] = None,
+        links: Optional[LinksInputType] = None,
+        tags: Optional[TagsInputType] = None,
+        terms: Optional[TermsInputType] = None,
+        domain: Optional[DomainInputType] = None,
+        training_jobs: Optional[Sequence[Union[str, DataProcessInstanceUrn]]] = None,
+        downstream_jobs: Optional[Sequence[Union[str, DataProcessInstanceUrn]]] = None,
+        extra_aspects: ExtraAspectsType = None,
+    ):
+        urn = MlModelGroupUrn(platform=platform, name=id, env=env)
+        super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
+
+        self._set_platform_instance(urn.platform, platform_instance)
+
+        # Set MLModelGroupProperties aspect
+        self._ensure_model_group_props(name=display_name or name)
+
+        if description is not None:
+            self.set_description(description)
+        if external_url is not None:
+            self.set_external_url(external_url)
+        if custom_properties is not None:
+            self.set_custom_properties(custom_properties)
+        if created is not None:
+            self.set_created(created)
+        if last_modified is not None:
+            self.set_last_modified(last_modified)
+
+        # Standard aspects
+        if owners is not None:
+            self.set_owners(owners)
+        if links is not None:
+            self.set_links(links)
+        if tags is not None:
+            self.set_tags(tags)
+        if terms is not None:
+            self.set_terms(terms)
+        if domain is not None:
+            self.set_domain(domain)
+
+        # ML model group specific aspects
+        if training_jobs is not None:
+            self.set_training_jobs(training_jobs)
+        if downstream_jobs is not None:
+            self.set_downstream_jobs(downstream_jobs)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: AspectBag) -> Self:
+        assert isinstance(urn, MlModelGroupUrn)
+        entity = cls(
+            platform=urn.platform,
+            id=urn.name,
+            env=urn.env,
+        )
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> MlModelGroupUrn:
+        return self._urn  # type: ignore
+
+    def _ensure_model_group_props(
+        self, *, name: Optional[str] = None
+    ) -> MLModelGroupPropertiesClass:
+        if name is not None:
+            return self._setdefault_aspect(MLModelGroupPropertiesClass(name=name))
+
+        props = self._get_aspect(MLModelGroupPropertiesClass)
+        if props is None:
+            # If we need properties but they don't exist and no name was provided
+            return self._setdefault_aspect(
+                MLModelGroupPropertiesClass(name=self.urn.name)
+            )
+        return props
+
+    @property
+    def name(self) -> Optional[str]:
+        return self._ensure_model_group_props().name
+
+    def set_name(self, display_name: str) -> None:
+        self._ensure_model_group_props().name = display_name
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_model_group_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_model_group_props().description = description
+
+    @property
+    def external_url(self) -> Optional[str]:
+        return self._ensure_model_group_props().externalUrl
+
+    def set_external_url(self, external_url: str) -> None:
+        self._ensure_model_group_props().externalUrl = external_url
+
+    @property
+    def custom_properties(self) -> Optional[Dict[str, str]]:
+        return self._ensure_model_group_props().customProperties
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
+        self._ensure_model_group_props().customProperties = custom_properties
+
+    @property
+    def created(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_group_props().created)
+
+    def set_created(self, created: datetime) -> None:
+        self._ensure_model_group_props().created = make_time_stamp(created)
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_group_props().lastModified)
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        self._ensure_model_group_props().lastModified = make_time_stamp(last_modified)
+
+    @property
+    def training_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_group_props().trainingJobs
+
+    def set_training_jobs(
+        self, training_jobs: Sequence[Union[str, DataProcessInstanceUrn]]
+    ) -> None:
+        self._ensure_model_group_props().trainingJobs = [
+            str(job) for job in training_jobs
+        ]
+
+    def add_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.trainingJobs is None:
+            props.trainingJobs = []
+        props.trainingJobs.append(str(training_job))
+
+    def remove_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.trainingJobs is not None:
+            props.trainingJobs = [
+                job for job in props.trainingJobs if job != str(training_job)
+            ]
+
+    @property
+    def downstream_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_group_props().downstreamJobs
+
+    def set_downstream_jobs(
+        self, downstream_jobs: Sequence[Union[str, DataProcessInstanceUrn]]
+    ) -> None:
+        self._ensure_model_group_props().downstreamJobs = [
+            str(job) for job in downstream_jobs
+        ]
+
+    def add_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.downstreamJobs is None:
+            props.downstreamJobs = []
+        props.downstreamJobs.append(str(downstream_job))
+
+    def remove_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_group_props()
+        if props.downstreamJobs is not None:
+            props.downstreamJobs = [
+                job for job in props.downstreamJobs if job != str(downstream_job)
+            ]

datahub/secret/datahub_secret_store.py

@@ -3,7 +3,8 @@ from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, validator
 
-from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.secret.datahub_secrets_client import DataHubSecretsClient
 from datahub.secret.secret_store import SecretStore
 

datahub/specific/dataset.py

@@ -292,3 +292,15 @@ class DatasetPatchBuilder(
                 value=timestamp,
             )
         return self
+
+    def set_external_url(
+        self, external_url: Optional[str] = None
+    ) -> "DatasetPatchBuilder":
+        if external_url is not None:
+            self._add_patch(
+                DatasetProperties.ASPECT_NAME,
+                "add",
+                path=("externalUrl",),
+                value=external_url,
+            )
+        return self

datahub/sql_parsing/fingerprint_utils.py

@@ -0,0 +1,6 @@
+import hashlib
+
+
+def generate_hash(text: str) -> str:
+    # Once we move to Python 3.9+, we can set `usedforsecurity=False`.
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -32,6 +32,7 @@ from datahub.metadata.urns import (
     SchemaFieldUrn,
     Urn,
 )
+from datahub.sql_parsing.fingerprint_utils import generate_hash
 from datahub.sql_parsing.schema_resolver import (
     SchemaResolver,
     SchemaResolverInterface,
@@ -49,7 +50,6 @@ from datahub.sql_parsing.sqlglot_lineage import (
 )
 from datahub.sql_parsing.sqlglot_utils import (
     _parse_statement,
-    generate_hash,
     get_query_fingerprint,
     try_format_query,
 )
@@ -155,6 +155,47 @@ class QueryMetadata:
             actor=(self.actor or _DEFAULT_USER_URN).urn(),
         )
 
+    def get_subjects(
+        self,
+        downstream_urn: Optional[str],
+        include_fields: bool,
+    ) -> List[UrnStr]:
+        query_subject_urns = OrderedSet[UrnStr]()
+        for upstream in self.upstreams:
+            query_subject_urns.add(upstream)
+            if include_fields:
+                for column in sorted(self.column_usage.get(upstream, [])):
+                    query_subject_urns.add(
+                        builder.make_schema_field_urn(upstream, column)
+                    )
+        if downstream_urn:
+            query_subject_urns.add(downstream_urn)
+            if include_fields:
+                for column_lineage in self.column_lineage:
+                    query_subject_urns.add(
+                        builder.make_schema_field_urn(
+                            downstream_urn, column_lineage.downstream.column
+                        )
+                    )
+        return list(query_subject_urns)
+
+    def make_query_properties(self) -> models.QueryPropertiesClass:
+        return models.QueryPropertiesClass(
+            statement=models.QueryStatementClass(
+                value=self.formatted_query_string,
+                language=models.QueryLanguageClass.SQL,
+            ),
+            source=models.QuerySourceClass.SYSTEM,
+            created=self.make_created_audit_stamp(),
+            lastModified=self.make_last_modified_audit_stamp(),
+        )
+
+
+def make_query_subjects(urns: List[UrnStr]) -> models.QuerySubjectsClass:
+    return models.QuerySubjectsClass(
+        subjects=[models.QuerySubjectClass(entity=urn) for urn in urns]
+    )
+
 
 @dataclasses.dataclass
 class KnownQueryLineageInfo:
@@ -1440,42 +1481,15 @@ class SqlParsingAggregator(Closeable):
             self.report.num_queries_skipped_due_to_filters += 1
             return
 
-        query_subject_urns = OrderedSet[UrnStr]()
-        for upstream in query.upstreams:
-            query_subject_urns.add(upstream)
-            if self.generate_query_subject_fields:
-                for column in sorted(query.column_usage.get(upstream, [])):
-                    query_subject_urns.add(
-                        builder.make_schema_field_urn(upstream, column)
-                    )
-        if downstream_urn:
-            query_subject_urns.add(downstream_urn)
-            if self.generate_query_subject_fields:
-                for column_lineage in query.column_lineage:
-                    query_subject_urns.add(
-                        builder.make_schema_field_urn(
-                            downstream_urn, column_lineage.downstream.column
-                        )
-                    )
-
         yield from MetadataChangeProposalWrapper.construct_many(
             entityUrn=self._query_urn(query_id),
             aspects=[
-                models.QueryPropertiesClass(
-                    statement=models.QueryStatementClass(
-                        value=query.formatted_query_string,
-                        language=models.QueryLanguageClass.SQL,
-                    ),
-                    source=models.QuerySourceClass.SYSTEM,
-                    created=query.make_created_audit_stamp(),
-                    lastModified=query.make_last_modified_audit_stamp(),
-                    origin=query.origin.urn() if query.origin else None,
-                ),
-                models.QuerySubjectsClass(
-                    subjects=[
-                        models.QuerySubjectClass(entity=urn)
-                        for urn in query_subject_urns
-                    ]
+                query.make_query_properties(),
+                make_query_subjects(
+                    query.get_subjects(
+                        downstream_urn=downstream_urn,
+                        include_fields=self.generate_query_subject_fields,
+                    )
                 ),
                 models.DataPlatformInstanceClass(
                     platform=self.platform.urn(),

datahub/sql_parsing/sqlglot_utils.py

@@ -1,7 +1,6 @@
 from datahub.sql_parsing._sqlglot_patch import SQLGLOT_PATCHED
 
 import functools
-import hashlib
 import logging
 import re
 from typing import Dict, Iterable, Optional, Tuple, Union
@@ -10,6 +9,8 @@ import sqlglot
 import sqlglot.errors
 import sqlglot.optimizer.eliminate_ctes
 
+from datahub.sql_parsing.fingerprint_utils import generate_hash
+
 assert SQLGLOT_PATCHED
 
 logger = logging.getLogger(__name__)
@@ -251,13 +252,11 @@ def generalize_query(expression: sqlglot.exp.ExpOrStr, dialect: DialectOrStr) ->
     return expression.transform(_strip_expression, copy=True).sql(dialect=dialect)
 
 
-def generate_hash(text: str) -> str:
-    # Once we move to Python 3.9+, we can set `usedforsecurity=False`.
-    return hashlib.sha256(text.encode("utf-8")).hexdigest()
-
-
 def get_query_fingerprint_debug(
-    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr, fast: bool = False
+    expression: sqlglot.exp.ExpOrStr,
+    platform: DialectOrStr,
+    fast: bool = False,
+    secondary_id: Optional[str] = None,
 ) -> Tuple[str, Optional[str]]:
     try:
         if not fast:
@@ -272,16 +271,18 @@ def get_query_fingerprint_debug(
         logger.debug("Failed to generalize query for fingerprinting: %s", e)
         expression_sql = None
 
-    fingerprint = generate_hash(
-        expression_sql
-        if expression_sql is not None
-        else _expression_to_string(expression, platform=platform)
-    )
+    text = expression_sql or _expression_to_string(expression, platform=platform)
+    if secondary_id:
+        text = text + " -- " + secondary_id
+    fingerprint = generate_hash(text=text)
     return fingerprint, expression_sql
 
 
 def get_query_fingerprint(
-    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr, fast: bool = False
+    expression: sqlglot.exp.ExpOrStr,
+    platform: DialectOrStr,
+    fast: bool = False,
+    secondary_id: Optional[str] = None,
 ) -> str:
     """Get a fingerprint for a SQL query.
 
@@ -298,12 +299,15 @@ def get_query_fingerprint(
     Args:
         expression: The SQL query to fingerprint.
         platform: The SQL dialect to use.
+        secondary_id: An optional additional id string to included in the final fingerprint.
 
     Returns:
         The fingerprint for the SQL query.
     """
 
-    return get_query_fingerprint_debug(expression, platform, fast=fast)[0]
+    return get_query_fingerprint_debug(
+        expression=expression, platform=platform, fast=fast, secondary_id=secondary_id
+    )[0]
 
 
 @functools.lru_cache(maxsize=FORMAT_QUERY_CACHE_SIZE)

datahub/telemetry/telemetry.py

@@ -352,10 +352,10 @@ class Telemetry:
             }
         else:
             return {
-                "server_type": server.server_config.get("datahub", {}).get(
+                "server_type": server.server_config.raw_config.get("datahub", {}).get(
                     "serverType", "missing"
                 ),
-                "server_version": server.server_config.get("versions", {})
+                "server_version": server.server_config.raw_config.get("versions", {})
                 .get("acryldata/datahub", {})
                 .get("version", "missing"),
                 "server_id": server.server_id or "missing",

datahub/testing/check_imports.py

@@ -9,7 +9,7 @@ def ensure_no_indirect_model_imports(dirs: List[pathlib.Path]) -> None:
     # If our needs become more complex, we should move to a proper linter.
     denied_imports = {
         "src.": "datahub.*",
-        "datahub.metadata._schema_classes": "datahub.metadata.schema_classes",
+        "datahub.metadata._internal_schema_classes": "datahub.metadata.schema_classes",
         "datahub.metadata._urns": "datahub.metadata.urns",
     }
     ignored_files = {

datahub/testing/mcp_diff.py

@@ -2,7 +2,7 @@ import dataclasses
 import json
 import re
 from collections import defaultdict
-from typing import Any, Dict, List, Sequence, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import deepdiff.serialization
 import yaml
@@ -34,6 +34,7 @@ class AspectForDiff:
     aspect_name: str
     aspect: Dict[str, Any] = dataclasses.field(hash=False)
     delta_info: "DeltaInfo" = dataclasses.field(hash=False, repr=False)
+    headers: Optional[Dict[str, str]] = dataclasses.field(default=None, hash=False)
 
     @classmethod
     def create_from_mcp(cls, idx: int, obj: Dict[str, Any]) -> "AspectForDiff":
@@ -44,6 +45,7 @@ class AspectForDiff:
             aspect_name=obj["aspectName"],
             aspect=aspect.get("json", aspect),
             delta_info=DeltaInfo(idx=idx, original=obj),
+            headers=obj.get("headers"),
         )
 
     def __repr__(self):
@@ -240,9 +242,12 @@ class MCPDiff:
                         s.append(serialize_aspect(ga.aspect))
                 for (i, old, new), diffs in aspect_diffs.aspects_changed.items():
                     s.append(self.report_aspect(old, i, "changed") + ":")
+
+                    print_aspects = False
                     for diff_level in diffs:
                         s.append(self.report_diff_level(diff_level, i))
-                    if verbose:
+                        print_aspects |= self.is_diff_level_on_aspect(diff_level)
+                    if verbose and print_aspects:
                         s.append(f"Old aspect:\n{serialize_aspect(old.aspect)}")
                         s.append(f"New aspect:\n{serialize_aspect(new.aspect)}")
 
@@ -271,6 +276,14 @@ class MCPDiff:
             f"root[{idx}].", ""
         )
 
+    @staticmethod
+    def is_diff_level_on_aspect(diff: DiffLevel) -> bool:
+        skip_print_fields = ["changeType", "headers"]
+        try:
+            return diff.path(output_format="list")[1] not in skip_print_fields
+        except IndexError:
+            return True
+
 
 def serialize_aspect(aspect: Union[AspectForDiff, Dict[str, Any]]) -> str:
     if isinstance(aspect, AspectForDiff):  # Unpack aspect

datahub/upgrade/upgrade.py

@@ -13,7 +13,9 @@ from pydantic import BaseModel
 from datahub._version import __version__
 from datahub.cli.config_utils import load_client_config
 from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.utilities.perf_timer import PerfTimer
+from datahub.utilities.server_config_util import RestServiceConfig
 
 log = logging.getLogger(__name__)
 
@@ -109,7 +111,7 @@ async def get_github_stats():
             return (latest_server_version, latest_server_date)
 
 
-async def get_server_config(gms_url: str, token: Optional[str]) -> dict:
+async def get_server_config(gms_url: str, token: Optional[str]) -> RestServiceConfig:
     import aiohttp
 
     headers = {
@@ -124,7 +126,7 @@ async def get_server_config(gms_url: str, token: Optional[str]) -> dict:
         config_endpoint = f"{gms_url}/config"
         async with session.get(config_endpoint, headers=headers) as dh_response:
             dh_response_json = await dh_response.json()
-            return dh_response_json
+            return RestServiceConfig(raw_config=dh_response_json)
 
 
 async def get_server_version_stats(
@@ -132,11 +134,12 @@ async def get_server_version_stats(
 ) -> Tuple[Optional[str], Optional[Version], Optional[datetime]]:
     import aiohttp
 
-    server_config = None
+    server_config: Optional[RestServiceConfig] = None
     if not server:
         try:
             # let's get the server from the cli config
             client_config = load_client_config()
+            client_config.client_mode = ClientMode.CLI
             host = client_config.server
             token = client_config.token
             server_config = await get_server_config(host, token)
@@ -150,15 +153,10 @@ async def get_server_version_stats(
     server_version: Optional[Version] = None
     current_server_release_date = None
     if server_config:
-        server_version_string = (
-            server_config.get("versions", {})
-            .get("acryldata/datahub", {})
-            .get("version")
-        )
-        commit_hash = (
-            server_config.get("versions", {}).get("acryldata/datahub", {}).get("commit")
-        )
-        server_type = server_config.get("datahub", {}).get("serverType", "unknown")
+        server_version_string = server_config.service_version
+        commit_hash = server_config.commit_hash
+        server_type = server_config.server_type
+
         if server_type == "quickstart" and commit_hash:
             async with aiohttp.ClientSession(
                 headers={"Accept": "application/vnd.github.v3+json"}

datahub/utilities/logging_manager.py

@@ -161,6 +161,7 @@ class _LogBuffer:
         self._buffer: Deque[str] = collections.deque(maxlen=maxlen)
 
     def write(self, line: str) -> None:
+        # We do not expect `line` to have a trailing newline.
         if len(line) > IN_MEMORY_LOG_BUFFER_MAX_LINE_LENGTH:
             line = line[:IN_MEMORY_LOG_BUFFER_MAX_LINE_LENGTH] + "[truncated]"
 
@@ -188,7 +189,13 @@ class _BufferLogHandler(logging.Handler):
             message = self.format(record)
         except TypeError as e:
             message = f"Error formatting log message: {e}\nMessage: {record.msg}, Args: {record.args}"
-        self._storage.write(message)
+
+        # For exception stack traces, the message is split over multiple lines,
+        # but we store it as a single string. Because we truncate based on line
+        # length, it's better for us to split it into multiple lines so that we
+        # don't lose any information on deeper stack traces.
+        for line in message.split("\n"):
+            self._storage.write(line)
 
 
 def _remove_all_handlers(logger: logging.Logger) -> None: