acryl-datahub 1.0.0.3rc12__py3-none-any.whl → 1.0.0.4rc2__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__)
@@ -55,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]:
         """
@@ -123,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")
 
@@ -135,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
@@ -146,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")
 
@@ -233,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:
@@ -243,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
@@ -254,7 +191,7 @@ class RestServiceConfig:
         Returns:
             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
@@ -277,18 +214,20 @@ class RestServiceConfig:
         # 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),
             ServiceFeature.DATAHUB_CLOUD: lambda: self.is_datahub_cloud,
         }
@@ -332,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:
         """
@@ -341,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: