acryl-datahub 1.3.1.1rc1__py3-none-any.whl → 1.3.1.1rc3__py3-none-any.whl

{acryl_datahub-1.3.1.1rc1.dist-info → acryl_datahub-1.3.1.1rc3.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.3.1.1rc1.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.3.1.1rc3.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=EeAHlBCqg0RG-oXfYfDsP4a2hg6rh7Jnhdkg-lWY3Es,323
+datahub/_version.py,sha256=m-yUiaArabPX6tKfFCNQY2lBFEsMvc1WrBPy2LEU7gw,323
 datahub/entrypoints.py,sha256=VcbU6Z47b_JKW1zI-WJMYIngm05FSogKLiuvFNtyNcI,9088
 datahub/errors.py,sha256=p5rFAdAGVCk4Lqolol1YvthceadUSwpaCxLXRcyCCFQ,676
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -323,7 +323,7 @@ datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py,sha256=MQk8BAHL
 datahub/ingestion/source/dremio/dremio_entities.py,sha256=e5NgsMMMiyNAEWeJX3iIasY-i8FL9bTozLyWCIFjrUU,16339
 datahub/ingestion/source/dremio/dremio_profiling.py,sha256=fxNlIHjZx2wrpU_zoGa89yvckgC99hj5Ki_Cnh3KRIw,12742
 datahub/ingestion/source/dremio/dremio_reporting.py,sha256=UEj-6FMdIWsry5535_kM2hLze5aPRMatTwvI0Bd2BSo,2140
-datahub/ingestion/source/dremio/dremio_source.py,sha256=iESzBbPLzSNBjNhkoeY0nlchgF13W8UthFc-3oCjAd8,25478
+datahub/ingestion/source/dremio/dremio_source.py,sha256=LANQSJwHhgYVVvitIH72vn1GOTARBYvroRofGPZuu00,25658
 datahub/ingestion/source/dremio/dremio_sql_queries.py,sha256=wA1hqKk9cKMJDyEdZRQcDDLZPGYwuNqrvleUHTkWgrQ,10508
 datahub/ingestion/source/dynamodb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/dynamodb/data_reader.py,sha256=vC77KpcP8LJN0g8wsPRDVw4sebv0ZWIP3tJkEIHaomA,3120
@@ -572,18 +572,19 @@ datahub/ingestion/source/tableau/tableau_server_wrapper.py,sha256=wsVD0SkGUwb-H9
 datahub/ingestion/source/tableau/tableau_validation.py,sha256=Hjbfc1AMIkGgzo5ffWXtNRjrxSxzHvw7-dYZDt4d3WE,1819
 datahub/ingestion/source/unity/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/unity/analyze_profiler.py,sha256=2pqkFY30CfN4aHgFZZntjeG0hNhBytZJvXC13VfTc1I,4689
-datahub/ingestion/source/unity/config.py,sha256=9vUpkuGOzUaatY3pf2U3Bl-GHkeBhbciUA2k3VYv50I,19912
-datahub/ingestion/source/unity/connection.py,sha256=iCsQhZ1vxzv1qQKTl_sFUZdmBLLIrNdu2X2V8hT7IGI,2441
-datahub/ingestion/source/unity/connection_test.py,sha256=Dwpz4AIc6ZDwq6pWmRCSCuDUgNjPP_bVAVJumgAAS4w,2661
+datahub/ingestion/source/unity/azure_auth_config.py,sha256=JZKrncqGvXlNP3e8ukYNPIoA2zwEDmjTtnGVc7brgvE,626
+datahub/ingestion/source/unity/config.py,sha256=IHlwcRplXf8BR3yPHounYy2Ua_5L2pen3eMhTCt6I_w,20586
+datahub/ingestion/source/unity/connection.py,sha256=K56rwh1LcCVKOmh3YO43Gn7tPburXNRWQQ1sl_gzgSA,2669
+datahub/ingestion/source/unity/connection_test.py,sha256=uAoDUOBPTFQ-1m0lQW3AZ2pXMVaIGkook3oR3BIQyn8,2683
 datahub/ingestion/source/unity/ge_profiler.py,sha256=NBRHZceq-f95iUn7u0h7cgcd9nAc48Aa-lmp_BqE0As,8409
 datahub/ingestion/source/unity/hive_metastore_proxy.py,sha256=IAWWJjaW0si_UF52Se2D7wmdYRY_afUG4QlVmQu6xaw,15351
 datahub/ingestion/source/unity/platform_resource_repository.py,sha256=znDmtddQB6B0owcWgOf5DACV2hbe7lseO7hwGszYTYs,541
-datahub/ingestion/source/unity/proxy.py,sha256=LclyOQ0rxlCuSaSIo5x8jW4PqMT19Jdg_fde1WbS1Gs,56606
+datahub/ingestion/source/unity/proxy.py,sha256=8IrS5wl1jQiAhFHJEp35F9EYZD5xK6b4MSSEcIovbhk,56898
 datahub/ingestion/source/unity/proxy_patch.py,sha256=gVYl5Fm_ase0iwBf3yDg7PE3bbTOl92RR-JgofHNkus,12374
 datahub/ingestion/source/unity/proxy_profiling.py,sha256=WLqvYP6MziaisA4LYL4T_GA-kPt6Xdde7bfaYsjYw40,9663
 datahub/ingestion/source/unity/proxy_types.py,sha256=3ChQa8tjrXd89o1twB988e_uiJSz0UWne_G1jcBUpJo,13702
 datahub/ingestion/source/unity/report.py,sha256=wa5ER1hLc-362iWS9MNwDJf_adfAb3tfge_O3wSAaTc,3291
-datahub/ingestion/source/unity/source.py,sha256=w3ub6ChKjaeKpRkvnBFRKQ5msm0WXjdJzM7387Nzjak,57405
+datahub/ingestion/source/unity/source.py,sha256=NQKgr5Vk1zXsR_UllHQWd1ATCWUM1SxWlZGYswHZqhU,57526
 datahub/ingestion/source/unity/tag_entities.py,sha256=-Z-XYc1XhquE-Eoksn9v0o11ZjV9CWz8n6zeXLbzluQ,7275
 datahub/ingestion/source/unity/usage.py,sha256=5CMQ0uOwuXfRq0rdxQbiHwrfSkvmQ5R2GNLqqZ2dOJw,13351
 datahub/ingestion/source/usage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -649,7 +650,7 @@ datahub/lite/lite_server.py,sha256=p9Oa2nNs65mqcssSIVOr7VOzWqfVstz6ZQEdT4f82S0,1
 datahub/lite/lite_util.py,sha256=DOIzuU78oxdyHGV-MjpS_iUzO_zF0oa9KJemFhfBp8U,4555
 datahub/metadata/__init__.py,sha256=AjhXPjI6cnpdcrBRrE5gOWo15vv2TTl2ctU4UAnUN7A,238
 datahub/metadata/_internal_schema_classes.py,sha256=PBMgJ92rA0JQlEKPcjcn_eFdjEVPZnapmi8gmHYgheQ,1086866
-datahub/metadata/schema.avsc,sha256=-P9qY8bKEj2ilPHKlYErL0eAxnkBALTIs7Rr0hzFRyg,787557
+datahub/metadata/schema.avsc,sha256=mj0wn3jv2XPajNcmXgBxuLdDeTvSYLWP877IJrvtf20,787709
 datahub/metadata/schema_classes.py,sha256=tPT8iHCak4IsZi_oL0nirbPpI8ETTPTZzapqLRpeKU4,1326
 datahub/metadata/urns.py,sha256=nfrCTExR-k2P9w272WVtWSN3xW1VUJngPwP3xnvULjU,1217
 datahub/metadata/_urns/__init__.py,sha256=cOF3GHMDgPhmbLKbN02NPpuLGHSu0qNgQyBRv08eqF0,243
@@ -933,7 +934,7 @@ datahub/metadata/schemas/PostKey.avsc,sha256=MEA-S-RonGrsrpwYEejR6CXRka2DcqP7S7q
 datahub/metadata/schemas/QuantitativeAnalyses.avsc,sha256=XrGYVoGUh7ZRjAbMORSiqF5Zl4a0IbYDPiQuSJhABZM,960
 datahub/metadata/schemas/QueryKey.avsc,sha256=VI4oIHvAO7f0lN_7V3QVuBfHcPz31c57XtW6IrlDfxc,518
 datahub/metadata/schemas/QueryProperties.avsc,sha256=26Q3zzuzJbUCUG7IJ3q_OEdcNbyzloZzDJWPbv3GQAk,5589
-datahub/metadata/schemas/QuerySubjects.avsc,sha256=PDM6Ek1AkjwLGDk17Rjg7KVtE2tbgFpc1XTGkhoG4MA,1437
+datahub/metadata/schemas/QuerySubjects.avsc,sha256=EpwpqwQBgTIcPlYDIR-J_4VlQ8Y66xSExMdxJxpctxI,1441
 datahub/metadata/schemas/QueryUsageStatistics.avsc,sha256=z1gfAnXdBoPEeERi5RESjrdBuS6AcIdqdN5JqWOSuNo,6192
 datahub/metadata/schemas/RelationshipChangeEvent.avsc,sha256=il7yiTXf9nkZvSyqH8GUu0j2jws_n-QlABXThLJ871c,8288
 datahub/metadata/schemas/RoleKey.avsc,sha256=Uas5jFViSHXhFqq8D4P6-UXqywOth3coztjQ5wA7wL0,449
@@ -948,7 +949,7 @@ datahub/metadata/schemas/SlackUserInfo.avsc,sha256=IY7InWaiDzJa3hJ9J4W3Eg8EUKuh0
 datahub/metadata/schemas/SourceCode.avsc,sha256=tUgo2rczO5x1fxw3fYNWQj-51vRNmNIj38b1wayA0aQ,1370
 datahub/metadata/schemas/Status.avsc,sha256=rPZSXSJdwnNywqNx2qll8cdt54aYgI-YUbRr3GK7h78,522
 datahub/metadata/schemas/StructuredProperties.avsc,sha256=qe45sKZ9XrLcf15Gt03Ttzt2J_kJYHvN-DAOSErSYuY,7028
-datahub/metadata/schemas/StructuredPropertyDefinition.avsc,sha256=TCvF80cW_2uEV2yJZOjAdxVV-Wr-4OsasSAM-5QDoAo,11921
+datahub/metadata/schemas/StructuredPropertyDefinition.avsc,sha256=n7WFG8ugTeBP3lB1KxrMPENw7aS73w8lkzRFcKcmNOk,12059
 datahub/metadata/schemas/StructuredPropertyKey.avsc,sha256=lp7tQBgeriEU1YMQ6a4-6aUGSWDqNl00lLDym97j1yI,618
 datahub/metadata/schemas/StructuredPropertySettings.avsc,sha256=oHPbDuV_Wr_zR9vmTwZZKRzeaBx5nyqN1zn0cfIRhuQ,4044
 datahub/metadata/schemas/SubTypes.avsc,sha256=bhXbzK020zDyQno97Xp05vmoMeZ82IGu2jz7pWDo3RQ,655
@@ -971,8 +972,8 @@ datahub/metadata/schemas/ViewProperties.avsc,sha256=3HhcbH5493dJUnEUtFMYMVfbYQ52
 datahub/metadata/schemas/__init__.py,sha256=kCcak_fBn_KyuysZTJIoipAzZ8EO44Amk4DWSEvplEY,581
 datahub/pydantic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/pydantic/compat.py,sha256=TUEo4kSEeOWVAhV6LQtst1phrpVgGtK4uif4OI5vQ2M,1937
-datahub/sdk/__init__.py,sha256=66OOcFi7qlnL6q72c_yUX2mWU2HudbOdRsC5CIoDxow,1922
-datahub/sdk/_all_entities.py,sha256=eQAmD_fcEHlTShe1_nHpdvHxLDN9njk9bdLnuTrYg8M,905
+datahub/sdk/__init__.py,sha256=Bo7pnilTfV4ucYy0zHHFQCymnNw7_dNuBnwPKowORC0,1954
+datahub/sdk/_all_entities.py,sha256=CllJAJ0gncg10Qm6VCSQTqDOU98P0Yt-gSh8BCZccec,946
 datahub/sdk/_attribution.py,sha256=0Trh8steVd27GOr9MKCZeawbuDD2_q3GIsZlCtHqEUg,1321
 datahub/sdk/_shared.py,sha256=tlyxyxgo7x-8sJhUG9VvUiLpxcOP-GIg4ToqPixagbw,34221
 datahub/sdk/_utils.py,sha256=oXE2BzsXE5zmSkCP3R1tObD4RHnPeH_ps83D_Dw9JaQ,1169
@@ -991,6 +992,7 @@ datahub/sdk/mlmodelgroup.py,sha256=wlZZHny0UORpF0fRYuVkWLSQwIHX_fWl5lPb1NKR6dM,8
 datahub/sdk/resolver_client.py,sha256=nKMAZJt2tRSGfKSzoREIh43PXqjM3umLiYkYHJjo1io,3243
 datahub/sdk/search_client.py,sha256=f2L_aOy-pPB9Mk7WdLSJ6Htp3OT4cEgDIrqnCweNtM8,3592
 datahub/sdk/search_filters.py,sha256=XG-V4V21DWwFDdBFowC0PsZ5I8H2snEfX_-0jFn5zYY,22929
+datahub/sdk/tag.py,sha256=jZHtL3hyiibrCXb0i_q8QsL8eT6dCT4egpqBWrTVe04,3123
 datahub/secret/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/secret/datahub_secret_store.py,sha256=XjmDElUUEzozw589-NN63LcTArquJhiEzxOKKTP-NcA,2665
 datahub/secret/datahub_secrets_client.py,sha256=nDmhziKdvseJHlaDVUcAwK8Fv8maeAaG-ktZtWG2b70,1316
@@ -1133,8 +1135,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.3.1.1rc1.dist-info/METADATA,sha256=rFziskaXXNthb4jDqBkMBVYPZXxgRIpVioW85ZFjUzE,184975
-acryl_datahub-1.3.1.1rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-acryl_datahub-1.3.1.1rc1.dist-info/entry_points.txt,sha256=pzsBoTx-D-iTcmpX8oCGCyzlHP2112EygUMzZWz56M8,10105
-acryl_datahub-1.3.1.1rc1.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.3.1.1rc1.dist-info/RECORD,,
+acryl_datahub-1.3.1.1rc3.dist-info/METADATA,sha256=r0p0khmFOz6-qenHWVttyw2qegGG8qriPU--__AHceA,184975
+acryl_datahub-1.3.1.1rc3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+acryl_datahub-1.3.1.1rc3.dist-info/entry_points.txt,sha256=pzsBoTx-D-iTcmpX8oCGCyzlHP2112EygUMzZWz56M8,10105
+acryl_datahub-1.3.1.1rc3.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.3.1.1rc3.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.3.1.1rc1"
+__version__ = "1.3.1.1rc3"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/source/dremio/dremio_source.py

@@ -336,10 +336,10 @@ class DremioSource(StatefulIngestionSourceBase):
             return
 
         dataset_urn = make_dataset_urn_with_platform_instance(
-            platform=self.get_platform(),
-            name=dataset_name,
-            platform_instance=self.config.platform_instance,
+            platform=make_data_platform_urn(self.get_platform()),
+            name=f"dremio.{dataset_name}",
             env=self.config.env,
+            platform_instance=self.config.platform_instance,
         )
 
         for dremio_mcp in self.dremio_aspects.populate_dataset_mcp(
@@ -419,10 +419,10 @@ class DremioSource(StatefulIngestionSourceBase):
         schema_str = ".".join(dataset_info.path)
         dataset_name = f"{schema_str}.{dataset_info.resource_name}".lower()
         dataset_urn = make_dataset_urn_with_platform_instance(
-            platform=self.get_platform(),
-            name=dataset_name,
-            platform_instance=self.config.platform_instance,
+            platform=make_data_platform_urn(self.get_platform()),
+            name=f"dremio.{dataset_name}",
             env=self.config.env,
+            platform_instance=self.config.platform_instance,
         )
         yield from self.profiler.get_workunits(dataset_info, dataset_urn)
 
@@ -434,10 +434,10 @@ class DremioSource(StatefulIngestionSourceBase):
         """
         upstream_urns = [
             make_dataset_urn_with_platform_instance(
-                platform=self.get_platform(),
-                name=upstream_table.lower(),
-                platform_instance=self.config.platform_instance,
+                platform=make_data_platform_urn(self.get_platform()),
+                name=f"dremio.{upstream_table.lower()}",
                 env=self.config.env,
+                platform_instance=self.config.platform_instance,
             )
             for upstream_table in parents
         ]
@@ -496,19 +496,19 @@ class DremioSource(StatefulIngestionSourceBase):
         if query.query and query.affected_dataset:
             upstream_urns = [
                 make_dataset_urn_with_platform_instance(
-                    platform=self.get_platform(),
-                    name=ds.lower(),
-                    platform_instance=self.config.platform_instance,
+                    platform=make_data_platform_urn(self.get_platform()),
+                    name=f"dremio.{ds.lower()}",
                     env=self.config.env,
+                    platform_instance=self.config.platform_instance,
                 )
                 for ds in query.queried_datasets
             ]
 
             downstream_urn = make_dataset_urn_with_platform_instance(
-                platform=self.get_platform(),
-                name=query.affected_dataset.lower(),
-                platform_instance=self.config.platform_instance,
+                platform=make_data_platform_urn(self.get_platform()),
+                name=f"dremio.{query.affected_dataset.lower()}",
                 env=self.config.env,
+                platform_instance=self.config.platform_instance,
             )
 
             # Add query to SqlParsingAggregator

datahub/ingestion/source/unity/azure_auth_config.py

@@ -0,0 +1,15 @@
+from pydantic import Field, SecretStr
+
+from datahub.configuration import ConfigModel
+
+
+class AzureAuthConfig(ConfigModel):
+    client_secret: SecretStr = Field(
+        description="Azure application client secret used for authentication. This is a confidential credential that should be kept secure."
+    )
+    client_id: str = Field(
+        description="Azure application (client) ID. This is the unique identifier for the registered Azure AD application.",
+    )
+    tenant_id: str = Field(
+        description="Azure tenant (directory) ID. This identifies the Azure AD tenant where the application is registered.",
+    )

datahub/ingestion/source/unity/config.py

@@ -413,6 +413,24 @@ class UnityCatalogSourceConfig(
             )
         return workspace_url
 
+    @model_validator(mode="before")
+    def either_token_or_azure_auth_provided(cls, values: dict) -> dict:
+        token = values.get("token")
+        azure_auth = values.get("azure_auth")
+
+        # Check if exactly one of the authentication methods is provided
+        if not token and not azure_auth:
+            raise ValueError(
+                "Either 'azure_auth' or 'token' (personal access token) must be provided in the configuration."
+            )
+
+        if token and azure_auth:
+            raise ValueError(
+                "Cannot specify both 'token' and 'azure_auth'. Please provide only one authentication method."
+            )
+
+        return values
+
     @field_validator("include_metastore", mode="after")
     @classmethod
     def include_metastore_warning(cls, v: bool) -> bool:

datahub/ingestion/source/unity/connection.py

@@ -8,6 +8,7 @@ from pydantic import Field
 
 from datahub.configuration.common import ConfigModel
 from datahub.ingestion.source.sql.sqlalchemy_uri import make_sqlalchemy_uri
+from datahub.ingestion.source.unity.azure_auth_config import AzureAuthConfig
 
 DATABRICKS = "databricks"
 
@@ -19,7 +20,12 @@ class UnityCatalogConnectionConfig(ConfigModel):
     """
 
     scheme: str = DATABRICKS
-    token: str = pydantic.Field(description="Databricks personal access token")
+    token: Optional[str] = pydantic.Field(
+        default=None, description="Databricks personal access token"
+    )
+    azure_auth: Optional[AzureAuthConfig] = Field(
+        default=None, description="Azure configuration"
+    )
     workspace_url: str = pydantic.Field(
         description="Databricks workspace url. e.g. https://my-workspace.cloud.databricks.com"
     )

datahub/ingestion/source/unity/connection_test.py

@@ -16,10 +16,10 @@ class UnityCatalogConnectionTest:
         self.report = UnityCatalogReport()
         self.proxy = UnityCatalogApiProxy(
             self.config.workspace_url,
-            self.config.token,
             self.config.profiling.warehouse_id,
             report=self.report,
             databricks_api_page_size=self.config.databricks_api_page_size,
+            personal_access_token=self.config.token,
         )
 
     def get_connection_test(self) -> TestConnectionReport:

datahub/ingestion/source/unity/proxy.py

@@ -44,6 +44,7 @@ from typing_extensions import assert_never
 from datahub._version import nice_version_name
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import parse_ts_millis
+from datahub.ingestion.source.unity.azure_auth_config import AzureAuthConfig
 from datahub.ingestion.source.unity.config import (
     LineageDataSource,
     UsageDataSource,
@@ -169,20 +170,31 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     def __init__(
         self,
         workspace_url: str,
-        personal_access_token: str,
         warehouse_id: Optional[str],
         report: UnityCatalogReport,
         hive_metastore_proxy: Optional[HiveMetastoreProxy] = None,
         lineage_data_source: LineageDataSource = LineageDataSource.AUTO,
         usage_data_source: UsageDataSource = UsageDataSource.AUTO,
         databricks_api_page_size: int = 0,
+        personal_access_token: Optional[str] = None,
+        azure_auth: Optional[AzureAuthConfig] = None,
     ):
-        self._workspace_client = WorkspaceClient(
-            host=workspace_url,
-            token=personal_access_token,
-            product="datahub",
-            product_version=nice_version_name(),
-        )
+        if azure_auth:
+            self._workspace_client = WorkspaceClient(
+                host=workspace_url,
+                azure_tenant_id=azure_auth.tenant_id,
+                azure_client_id=azure_auth.client_id,
+                azure_client_secret=azure_auth.client_secret.get_secret_value(),
+                product="datahub",
+                product_version=nice_version_name(),
+            )
+        else:
+            self._workspace_client = WorkspaceClient(
+                host=workspace_url,
+                token=personal_access_token,
+                product="datahub",
+                product_version=nice_version_name(),
+            )
         self.warehouse_id = warehouse_id or ""
         self.report = report
         self.hive_metastore_proxy = hive_metastore_proxy
@@ -336,24 +348,18 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                             except (json.JSONDecodeError, TypeError) as e:
                                 logger.debug(f"Failed to parse outputs JSON: {e}")
 
-                        if "parameters" in signature_raw:
+                        if "params" in signature_raw:
                             try:
-                                signature_data["parameters"] = json.loads(
-                                    signature_raw["parameters"]
+                                signature_data["params"] = json.loads(
+                                    signature_raw["params"]
                                 )
                             except (json.JSONDecodeError, TypeError) as e:
-                                logger.debug(f"Failed to parse parameters JSON: {e}")
+                                logger.debug(f"Failed to parse params JSON: {e}")
 
                         return ModelSignature(
-                            inputs=signature_data["inputs"]
-                            if "inputs" in signature_raw
-                            else None,
-                            outputs=signature_data["outputs"]
-                            if "outputs" in signature_raw
-                            else None,
-                            parameters=signature_data["parameters"]
-                            if "parameters" in signature_raw
-                            else None,
+                            inputs=signature_data.get("inputs"),
+                            outputs=signature_data.get("outputs"),
+                            parameters=signature_data.get("params"),
                         )
                     else:
                         logger.debug(

datahub/ingestion/source/unity/source.py

@@ -211,13 +211,14 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
 
         self.unity_catalog_api_proxy = UnityCatalogApiProxy(
             config.workspace_url,
-            config.token,
             config.warehouse_id,
             report=self.report,
             hive_metastore_proxy=self.hive_metastore_proxy,
             lineage_data_source=config.lineage_data_source,
             usage_data_source=config.usage_data_source,
             databricks_api_page_size=config.databricks_api_page_size,
+            personal_access_token=config.token if config.token else None,
+            azure_auth=config.azure_auth if config.azure_auth else None,
         )
 
         self.external_url_base = urljoin(self.config.workspace_url, "/explore/data")

datahub/metadata/schema.avsc

@@ -4824,7 +4824,7 @@
               {
                 "Searchable": {
                   "fieldName": "entities",
-                  "fieldType": "URN"
+                  "fieldType": "KEYWORD"
                 },
                 "java": {
                   "class": "com.linkedin.pegasus2avro.common.urn.Urn"
@@ -20303,7 +20303,12 @@
         "doc": "The fully qualified name of the property. e.g. io.acryl.datahub.myProperty"
       },
       {
-        "Searchable": {},
+        "Searchable": {
+          "enableAutocomplete": true,
+          "fieldType": "WORD_GRAM",
+          "searchLabel": "entityName",
+          "searchTier": 1
+        },
         "type": [
           "null",
           "string"

datahub/metadata/schemas/QuerySubjects.avsc

@@ -17,7 +17,7 @@
             {
               "Searchable": {
                 "fieldName": "entities",
-                "fieldType": "URN"
+                "fieldType": "KEYWORD"
               },
               "java": {
                 "class": "com.linkedin.pegasus2avro.common.urn.Urn"

datahub/metadata/schemas/StructuredPropertyDefinition.avsc

@@ -17,7 +17,12 @@
       "doc": "The fully qualified name of the property. e.g. io.acryl.datahub.myProperty"
     },
     {
-      "Searchable": {},
+      "Searchable": {
+        "enableAutocomplete": true,
+        "fieldType": "WORD_GRAM",
+        "searchLabel": "entityName",
+        "searchTier": 1
+      },
       "type": [
         "null",
         "string"

datahub/sdk/__init__.py

@@ -28,6 +28,7 @@ from datahub.sdk.main_client import DataHubClient
 from datahub.sdk.mlmodel import MLModel
 from datahub.sdk.mlmodelgroup import MLModelGroup
 from datahub.sdk.search_filters import Filter, FilterDsl
+from datahub.sdk.tag import Tag
 
 # We want to print out the warning if people do `from datahub.sdk import X`.
 # But we don't want to print out warnings if they're doing a more direct

datahub/sdk/_all_entities.py

@@ -9,6 +9,7 @@ from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
 from datahub.sdk.mlmodel import MLModel
 from datahub.sdk.mlmodelgroup import MLModelGroup
+from datahub.sdk.tag import Tag
 
 # Base entity classes that don't have circular dependencies
 # Those that do are imported in the EntityClient where needed
@@ -22,6 +23,7 @@ ENTITY_CLASSES_LIST: List[Type[Entity]] = [
     DataJob,
     Dashboard,
     Chart,
+    Tag,
 ]
 
 # Create the mapping of entity types to classes

datahub/sdk/tag.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from typing import Optional, Type
+
+from typing_extensions import Self
+
+import datahub.metadata.schema_classes as models
+from datahub.metadata.urns import TagUrn, Urn
+from datahub.sdk._shared import (
+    HasOwnership,
+    OwnersInputType,
+)
+from datahub.sdk.entity import Entity, ExtraAspectsType
+
+
+class Tag(
+    HasOwnership,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[TagUrn]:
+        return TagUrn
+
+    def __init__(
+        self,
+        *,
+        # Identity.
+        name: str,
+        # Tag properties.
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        color: Optional[str] = None,
+        # Standard aspects.
+        owners: Optional[OwnersInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
+    ):
+        """Initialize a new Tag instance."""
+        urn = TagUrn(name=name)
+        super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
+
+        self._ensure_tag_props(
+            display_name=display_name or name,
+            description=description,
+            color=color,
+        )
+
+        if owners is not None:
+            self.set_owners(owners)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
+        assert isinstance(urn, TagUrn)
+        entity = cls(name=urn.name)
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> TagUrn:
+        assert isinstance(self._urn, TagUrn)
+        return self._urn
+
+    def _ensure_tag_props(
+        self,
+        *,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        color: Optional[str] = None,
+    ) -> models.TagPropertiesClass:
+        existing_props = self._get_aspect(models.TagPropertiesClass)
+        if existing_props is not None:
+            if display_name is not None:
+                existing_props.name = display_name
+            if description is not None:
+                existing_props.description = description
+            if color is not None:
+                existing_props.colorHex = color
+            return existing_props
+
+        return self._setdefault_aspect(
+            models.TagPropertiesClass(
+                name=display_name or self.urn.name,
+                description=description,
+                colorHex=color,
+            )
+        )
+
+    @property
+    def name(self) -> str:
+        return self.urn.name
+
+    @property
+    def display_name(self) -> str:
+        return self._ensure_tag_props().name
+
+    def set_display_name(self, display_name: str) -> None:
+        self._ensure_tag_props(display_name=display_name)
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_tag_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_tag_props(description=description)
+
+    @property
+    def color(self) -> Optional[str]:
+        return self._ensure_tag_props().colorHex
+
+    def set_color(self, color: str) -> None:
+        self._ensure_tag_props(color=color)