acryl-datahub 1.0.0.3rc11__py3-none-any.whl → 1.0.0.4rc1__py3-none-any.whl

datahub/metadata/schemas/Operation.avsc

@@ -254,6 +254,23 @@
       "type": "long",
       "name": "lastUpdatedTimestamp",
       "doc": "The time at which the operation occurred. Would be better named 'operationTime'"
+    },
+    {
+      "TimeseriesFieldCollection": {
+        "key": "query"
+      },
+      "type": [
+        "null",
+        {
+          "type": "array",
+          "items": "string"
+        }
+      ],
+      "name": "queries",
+      "default": null,
+      "doc": "Which queries were used in this operation.",
+      "Urn": "Urn",
+      "urn_is_array": true
     }
   ],
   "doc": "Operational info for an entity."

datahub/sdk/main_client.py

@@ -10,6 +10,13 @@ from datahub.sdk.lineage_client import LineageClient
 from datahub.sdk.resolver_client import ResolverClient
 from datahub.sdk.search_client import SearchClient
 
+try:
+    from acryl_datahub_cloud._sdk_extras import (  # type: ignore[import-not-found]
+        AssertionClient,
+    )
+except ImportError:
+    AssertionClient = None
+
 
 class DataHubClient:
     """Main client for interacting with DataHub.
@@ -103,3 +110,11 @@ class DataHubClient:
     @property
     def lineage(self) -> LineageClient:
         return LineageClient(self)
+
+    @property
+    def assertion(self) -> AssertionClient:  # type: ignore[return-value]  # Type is not available if assertion_client is not installed
+        if AssertionClient is None:
+            raise SdkUsageError(
+                "AssertionClient is not installed, please install it with `pip install acryl-datahub-cloud`"
+            )
+        return AssertionClient(self)

datahub/sql_parsing/_sqlglot_patch.py

@@ -163,8 +163,7 @@ def _patch_lineage() -> None:
 -    source_columns = set(find_all_in_scope(select, exp.Column))
 +    source_columns = list(find_all_in_scope(select, exp.Column))
 
--    # If the source is a UDTF find columns used in the UTDF to generate the table
-+    # If the source is a UDTF find columns used in the UDTF to generate the table
+     # If the source is a UDTF find columns used in the UDTF to generate the table
 +    source = scope.expression
      if isinstance(source, exp.UDTF):
 -        source_columns |= set(source.find_all(exp.Column))

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -1753,8 +1753,9 @@ class SqlParsingAggregator(Closeable):
             operationType=operation_type,
             lastUpdatedTimestamp=make_ts_millis(query.latest_timestamp),
             actor=query.actor.urn() if query.actor else None,
-            customProperties=(
-                {"query_urn": self._query_urn(query_id)}
+            sourceType=models.OperationSourceTypeClass.DATA_PLATFORM,
+            queries=(
+                [self._query_urn(query_id)]
                 if self.can_generate_query(query_id)
                 else None
             ),

datahub/utilities/server_config_util.py

@@ -10,11 +10,6 @@ from typing import (
     Union,
 )
 
-import requests
-
-from datahub.configuration.common import (
-    ConfigurationError,
-)
 from datahub.telemetry.telemetry import suppress_telemetry
 
 logger = logging.getLogger(__name__)
@@ -44,15 +39,8 @@ class ServiceFeature(Enum):
 
 
 _REQUIRED_VERSION_OPENAPI_TRACING = {
-    "acryl": (
-        0,
-        3,
-        11,
-        0,
-    ),  # Requires v0.3.11.0 or higher for acryl versions
-    "cloud": (0, 3, 11, 0),  # Special case for '-cloud' suffix
-    "any_suffix": (0, 3, 11, 0),  # Generic requirement for any other suffix
-    "none": (1, 0, 1, 0),  # Requirement for versions without suffix
+    "cloud": (0, 3, 11, 0),
+    "core": (1, 0, 1, 0),
 }
 
 
@@ -62,66 +50,9 @@ class RestServiceConfig:
     A class to represent REST service configuration with semantic version parsing capabilities.
     """
 
-    session: Optional[requests.Session] = None
-    url: Optional[str] = None
     raw_config: Dict[str, Any] = field(default_factory=dict)
     _version_cache: Optional[Tuple[int, int, int, int]] = None
 
-    def fetch_config(self) -> Dict[str, Any]:
-        """
-        Fetch configuration from the server if not already loaded.
-
-        Returns:
-            The configuration dictionary
-
-        Raises:
-            ConfigurationError: If there's an error fetching or validating the configuration
-        """
-        if not self.raw_config:
-            if self.session is None or self.url is None:
-                raise ConfigurationError(
-                    "Session and URL are required to load configuration"
-                )
-
-            response = self.session.get(self.url)
-
-            if response.status_code == 200:
-                config = response.json()
-
-                # Validate that we're connected to the correct service
-                if config.get("noCode") == "true":
-                    self.raw_config = config
-                else:
-                    raise ConfigurationError(
-                        "You seem to have connected to the frontend service instead of the GMS endpoint. "
-                        "The rest emitter should connect to DataHub GMS (usually <datahub-gms-host>:8080) or Frontend GMS API (usually <frontend>:9002/api/gms). "
-                        "For Acryl users, the endpoint should be https://<name>.acryl.io/gms"
-                    )
-            else:
-                logger.debug(
-                    f"Unable to connect to {self.url} with status_code: {response.status_code}. Response: {response.text}"
-                )
-
-                if response.status_code == 401:
-                    message = f"Unable to connect to {self.url} - got an authentication error: {response.text}."
-                else:
-                    message = f"Unable to connect to {self.url} with status_code: {response.status_code}."
-
-                message += "\nPlease check your configuration and make sure you are talking to the DataHub GMS (usually <datahub-gms-host>:8080) or Frontend GMS API (usually <frontend>:9002/api/gms)."
-                raise ConfigurationError(message)
-
-        return self.raw_config
-
-    @property
-    def config(self) -> Dict[str, Any]:
-        """
-        Get the full configuration dictionary, loading it if necessary.
-
-        Returns:
-            The configuration dictionary
-        """
-        return self.fetch_config()
-
     @property
     def commit_hash(self) -> Optional[str]:
         """
@@ -130,7 +61,7 @@ class RestServiceConfig:
         Returns:
             The commit hash or None if not found
         """
-        versions = self.config.get("versions") or {}
+        versions = self.raw_config.get("versions") or {}
         datahub_info = versions.get("acryldata/datahub") or {}
         return datahub_info.get("commit")
 
@@ -142,7 +73,7 @@ class RestServiceConfig:
         Returns:
             The server type or "unknown" if not found
         """
-        datahub = self.config.get("datahub") or {}
+        datahub = self.raw_config.get("datahub") or {}
         return datahub.get("serverType", "unknown")
 
     @property
@@ -153,8 +84,7 @@ class RestServiceConfig:
         Returns:
             The version string or None if not found
         """
-        config = self.fetch_config()
-        versions = config.get("versions") or {}
+        versions = self.raw_config.get("versions") or {}
         datahub_info = versions.get("acryldata/datahub") or {}
         return datahub_info.get("version")
 
@@ -240,7 +170,7 @@ class RestServiceConfig:
         Returns:
             True if noCode is set to "true"
         """
-        return self.config.get("noCode") == "true"
+        return self.raw_config.get("noCode") == "true"
 
     @property
     def is_managed_ingestion_enabled(self) -> bool:
@@ -250,7 +180,7 @@ class RestServiceConfig:
         Returns:
             True if managedIngestion.enabled is True
         """
-        managed_ingestion = self.config.get("managedIngestion") or {}
+        managed_ingestion = self.raw_config.get("managedIngestion") or {}
         return managed_ingestion.get("enabled", False)
 
     @property
@@ -259,26 +189,21 @@ class RestServiceConfig:
         Check if DataHub Cloud is enabled.
 
         Returns:
-            True if the server environment is not 'oss'
+            True if the server environment is not 'core'
         """
-        datahub_config = self.config.get("datahub") or {}
+        datahub_config = self.raw_config.get("datahub") or {}
         server_env = datahub_config.get("serverEnv")
 
         # Return False if serverEnv is None or empty string
         if not server_env:
             return False
 
-        return server_env != "oss"
+        return server_env != "core"
 
     def supports_feature(self, feature: ServiceFeature) -> bool:
         """
-        Determines whether a specific feature is supported based on service version.
-
-        Version categorization follows these rules:
-        1. Has '-acryl' suffix (highest priority)
-        2. Has a specific known suffix (e.g. '-other')
-        3. Has some other suffix (catchall for any suffix)
-        4. No suffix
+        Determines whether a specific feature is supported based on service version
+        and whether this is a cloud deployment or not.
 
         Args:
             feature: Feature enum value to check
@@ -286,68 +211,54 @@ class RestServiceConfig:
         Returns:
             Boolean indicating whether the feature is supported
         """
-        version = self.service_version
-        if not version:
-            return False
-
-        # Determine the suffix category
-        suffix_category = "none"  # Default: no suffix
-
-        if "-" in version:
-            suffix = version.split("-", 1)[1]
-
-            if suffix == "acryl":
-                suffix_category = "acryl"
-            elif suffix == "cloud":  # Example of a specific override
-                suffix_category = "cloud"
-            else:
-                suffix_category = "any_suffix"  # Catchall for any other suffix
-
-        # Define feature requirements based on version scheme
-        # This can be expanded to include more features
-        feature_requirements = {
-            ServiceFeature.OPEN_API_SDK: _REQUIRED_VERSION_OPENAPI_TRACING,
-            ServiceFeature.API_TRACING: _REQUIRED_VERSION_OPENAPI_TRACING,
-            # Additional features can be defined here
-        }
-
-        # Special handling for features that rely on config flags instead of version
+        # Special handling for features that rely on config flags
         config_based_features = {
             ServiceFeature.NO_CODE: lambda: self.is_no_code_enabled,
-            ServiceFeature.STATEFUL_INGESTION: lambda: self.config.get(
+            ServiceFeature.STATEFUL_INGESTION: lambda: self.raw_config.get(
                 "statefulIngestionCapable", False
             )
             is True,
-            ServiceFeature.IMPACT_ANALYSIS: lambda: self.config.get(
+            ServiceFeature.IMPACT_ANALYSIS: lambda: self.raw_config.get(
                 "supportsImpactAnalysis", False
             )
             is True,
-            ServiceFeature.PATCH_CAPABLE: lambda: self.config.get("patchCapable", False)
+            ServiceFeature.PATCH_CAPABLE: lambda: self.raw_config.get(
+                "patchCapable", False
+            )
             is True,
             ServiceFeature.CLI_TELEMETRY: lambda: (
-                self.config.get("telemetry") or {}
+                self.raw_config.get("telemetry") or {}
             ).get("enabledCli", None),
-            # Add more config-based feature checks as needed
+            ServiceFeature.DATAHUB_CLOUD: lambda: self.is_datahub_cloud,
         }
 
         # Check if this is a config-based feature
         if feature in config_based_features:
             return config_based_features[feature]()
 
+        # For environment-based features, determine requirements based on cloud vs. non-cloud
+        deployment_type = "cloud" if self.is_datahub_cloud else "core"
+
+        # Define feature requirements
+        feature_requirements = {
+            ServiceFeature.OPEN_API_SDK: _REQUIRED_VERSION_OPENAPI_TRACING,
+            ServiceFeature.API_TRACING: _REQUIRED_VERSION_OPENAPI_TRACING,
+            # Additional features can be defined here
+        }
+
         # Check if the feature exists in our requirements dictionary
         if feature not in feature_requirements:
             # Unknown feature, assume not supported
             return False
 
-        # Get version requirements for this feature and version category
+        # Get version requirements for this feature and deployment type
         feature_reqs = feature_requirements[feature]
-        requirements = feature_reqs.get(suffix_category)
+        requirements = feature_reqs.get(deployment_type)
 
         if not requirements:
-            # Fallback to the no-suffix requirements if specific requirements aren't defined
-            requirements = feature_reqs.get(
-                "none", (99, 99, 99, 99)
-            )  # Very high version if none defined
+            # If no specific requirements defined for this deployment type,
+            # assume feature is not supported
+            return False
 
         # Check if the current version meets the requirements
         req_major, req_minor, req_patch, req_build = requirements
@@ -360,7 +271,7 @@ class RestServiceConfig:
         Returns:
             A string representation of the configuration dictionary
         """
-        return str(self.config)
+        return str(self.raw_config)
 
     def __repr__(self) -> str:
         """
@@ -369,7 +280,7 @@ class RestServiceConfig:
         Returns:
             A string representation that can be used with pprint
         """
-        return str(self.config)
+        return str(self.raw_config)
 
 
 def set_gms_config(config: Union[Dict[str, Any], RestServiceConfig]) -> None: