acryl-datahub 0.15.0.3__py3-none-any.whl → 0.15.0.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of acryl-datahub might be problematic. Click here for more details.

Files changed (36) hide show
  1. acryl_datahub-0.15.0.4.dist-info/LICENSE +202 -0
  2. {acryl_datahub-0.15.0.3.dist-info → acryl_datahub-0.15.0.4.dist-info}/METADATA +2417 -2414
  3. {acryl_datahub-0.15.0.3.dist-info → acryl_datahub-0.15.0.4.dist-info}/RECORD +36 -33
  4. datahub/__init__.py +1 -1
  5. datahub/cli/container_cli.py +108 -0
  6. datahub/emitter/enum_helpers.py +4 -2
  7. datahub/emitter/mce_builder.py +4 -0
  8. datahub/emitter/mcp_builder.py +19 -0
  9. datahub/entrypoints.py +2 -0
  10. datahub/ingestion/api/decorators.py +2 -0
  11. datahub/ingestion/api/registry.py +3 -1
  12. datahub/ingestion/api/sink.py +12 -0
  13. datahub/ingestion/api/source.py +5 -2
  14. datahub/ingestion/source/aws/glue.py +11 -5
  15. datahub/ingestion/source/aws/s3_util.py +1 -24
  16. datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py +2 -2
  17. datahub/ingestion/source/dbt/dbt_common.py +2 -2
  18. datahub/ingestion/source/powerbi/powerbi.py +4 -4
  19. datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py +6 -6
  20. datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py +24 -18
  21. datahub/ingestion/source/s3/source.py +6 -2
  22. datahub/ingestion/source/slack/slack.py +6 -0
  23. datahub/ingestion/source/sql/hive_metastore.py +3 -3
  24. datahub/ingestion/source/sql/mssql/job_models.py +2 -2
  25. datahub/ingestion/source/sql/mssql/source.py +26 -11
  26. datahub/ingestion/source/sql/teradata.py +2 -2
  27. datahub/ingestion/source/tableau/tableau.py +23 -10
  28. datahub/metadata/_schema_classes.py +401 -401
  29. datahub/metadata/_urns/urn_defs.py +1857 -1408
  30. datahub/metadata/schema.avsc +16624 -16266
  31. datahub/sql_parsing/sql_parsing_aggregator.py +3 -3
  32. datahub/utilities/groupby.py +17 -0
  33. datahub/utilities/urns/_urn_base.py +6 -2
  34. {acryl_datahub-0.15.0.3.dist-info → acryl_datahub-0.15.0.4.dist-info}/WHEEL +0 -0
  35. {acryl_datahub-0.15.0.3.dist-info → acryl_datahub-0.15.0.4.dist-info}/entry_points.txt +0 -0
  36. {acryl_datahub-0.15.0.3.dist-info → acryl_datahub-0.15.0.4.dist-info}/top_level.txt +0 -0
@@ -28,14 +28,25 @@ class StructuredPropertyUrn(_SpecificUrn):
28
28
  ENTITY_TYPE: ClassVar[str] = "structuredProperty"
29
29
  _URN_PARTS: ClassVar[int] = 1
30
30
 
31
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
31
+ def __init__(self, id: Union["StructuredPropertyUrn", str], *, _allow_coercion: bool = True) -> None:
32
32
  if _allow_coercion:
33
33
  # Field coercion logic (if any is required).
34
- id = UrnEncoder.encode_string(id)
34
+ if isinstance(id, str):
35
+ if id.startswith('urn:li:'):
36
+ try:
37
+ id = StructuredPropertyUrn.from_string(id)
38
+ except InvalidUrnError:
39
+ raise InvalidUrnError(f'Expecting a StructuredPropertyUrn but got {id}')
40
+ else:
41
+ id = UrnEncoder.encode_string(id)
35
42
 
36
43
  # Validation logic.
37
44
  if not id:
38
45
  raise InvalidUrnError("StructuredPropertyUrn id cannot be empty")
46
+ if isinstance(id, StructuredPropertyUrn):
47
+ id = id.id
48
+ elif isinstance(id, Urn):
49
+ raise InvalidUrnError(f'Expecting a StructuredPropertyUrn but got {id}')
39
50
  if UrnEncoder.contains_reserved_char(id):
40
51
  raise InvalidUrnError(f'StructuredPropertyUrn id contains reserved characters')
41
52
 
@@ -67,44 +78,55 @@ class StructuredPropertyUrn(_SpecificUrn):
67
78
  return self.entity_ids[0]
68
79
 
69
80
  if TYPE_CHECKING:
70
- from datahub.metadata.schema_classes import DataTypeKeyClass
81
+ from datahub.metadata.schema_classes import BusinessAttributeKeyClass
71
82
 
72
- class DataTypeUrn(_SpecificUrn):
73
- ENTITY_TYPE: ClassVar[str] = "dataType"
83
+ class BusinessAttributeUrn(_SpecificUrn):
84
+ ENTITY_TYPE: ClassVar[str] = "businessAttribute"
74
85
  _URN_PARTS: ClassVar[int] = 1
75
86
 
76
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
87
+ def __init__(self, id: Union["BusinessAttributeUrn", str], *, _allow_coercion: bool = True) -> None:
77
88
  if _allow_coercion:
78
89
  # Field coercion logic (if any is required).
79
- id = UrnEncoder.encode_string(id)
90
+ if isinstance(id, str):
91
+ if id.startswith('urn:li:'):
92
+ try:
93
+ id = BusinessAttributeUrn.from_string(id)
94
+ except InvalidUrnError:
95
+ raise InvalidUrnError(f'Expecting a BusinessAttributeUrn but got {id}')
96
+ else:
97
+ id = UrnEncoder.encode_string(id)
80
98
 
81
99
  # Validation logic.
82
100
  if not id:
83
- raise InvalidUrnError("DataTypeUrn id cannot be empty")
101
+ raise InvalidUrnError("BusinessAttributeUrn id cannot be empty")
102
+ if isinstance(id, BusinessAttributeUrn):
103
+ id = id.id
104
+ elif isinstance(id, Urn):
105
+ raise InvalidUrnError(f'Expecting a BusinessAttributeUrn but got {id}')
84
106
  if UrnEncoder.contains_reserved_char(id):
85
- raise InvalidUrnError(f'DataTypeUrn id contains reserved characters')
107
+ raise InvalidUrnError(f'BusinessAttributeUrn id contains reserved characters')
86
108
 
87
109
  super().__init__(self.ENTITY_TYPE, [id])
88
110
 
89
111
  @classmethod
90
- def _parse_ids(cls, entity_ids: List[str]) -> "DataTypeUrn":
112
+ def _parse_ids(cls, entity_ids: List[str]) -> "BusinessAttributeUrn":
91
113
  if len(entity_ids) != cls._URN_PARTS:
92
- raise InvalidUrnError(f"DataTypeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
114
+ raise InvalidUrnError(f"BusinessAttributeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
93
115
  return cls(id=entity_ids[0], _allow_coercion=False)
94
116
 
95
117
  @classmethod
96
- def underlying_key_aspect_type(cls) -> Type["DataTypeKeyClass"]:
97
- from datahub.metadata.schema_classes import DataTypeKeyClass
118
+ def underlying_key_aspect_type(cls) -> Type["BusinessAttributeKeyClass"]:
119
+ from datahub.metadata.schema_classes import BusinessAttributeKeyClass
98
120
 
99
- return DataTypeKeyClass
121
+ return BusinessAttributeKeyClass
100
122
 
101
- def to_key_aspect(self) -> "DataTypeKeyClass":
102
- from datahub.metadata.schema_classes import DataTypeKeyClass
123
+ def to_key_aspect(self) -> "BusinessAttributeKeyClass":
124
+ from datahub.metadata.schema_classes import BusinessAttributeKeyClass
103
125
 
104
- return DataTypeKeyClass(id=self.id)
126
+ return BusinessAttributeKeyClass(id=self.id)
105
127
 
106
128
  @classmethod
107
- def from_key_aspect(cls, key_aspect: "DataTypeKeyClass") -> "DataTypeUrn":
129
+ def from_key_aspect(cls, key_aspect: "BusinessAttributeKeyClass") -> "BusinessAttributeUrn":
108
130
  return cls(id=key_aspect.id)
109
131
 
110
132
  @property
@@ -112,94 +134,167 @@ class DataTypeUrn(_SpecificUrn):
112
134
  return self.entity_ids[0]
113
135
 
114
136
  if TYPE_CHECKING:
115
- from datahub.metadata.schema_classes import CorpGroupKeyClass
137
+ from datahub.metadata.schema_classes import DataProductKeyClass
116
138
 
117
- class CorpGroupUrn(_SpecificUrn):
118
- ENTITY_TYPE: ClassVar[str] = "corpGroup"
139
+ class DataProductUrn(_SpecificUrn):
140
+ ENTITY_TYPE: ClassVar[str] = "dataProduct"
119
141
  _URN_PARTS: ClassVar[int] = 1
120
142
 
121
- def __init__(self, name: str, *, _allow_coercion: bool = True) -> None:
143
+ def __init__(self, id: Union["DataProductUrn", str], *, _allow_coercion: bool = True) -> None:
122
144
  if _allow_coercion:
123
145
  # Field coercion logic (if any is required).
124
- name = UrnEncoder.encode_string(name)
146
+ if isinstance(id, str):
147
+ if id.startswith('urn:li:'):
148
+ try:
149
+ id = DataProductUrn.from_string(id)
150
+ except InvalidUrnError:
151
+ raise InvalidUrnError(f'Expecting a DataProductUrn but got {id}')
152
+ else:
153
+ id = UrnEncoder.encode_string(id)
125
154
 
126
155
  # Validation logic.
127
- if not name:
128
- raise InvalidUrnError("CorpGroupUrn name cannot be empty")
129
- if UrnEncoder.contains_reserved_char(name):
130
- raise InvalidUrnError(f'CorpGroupUrn name contains reserved characters')
156
+ if not id:
157
+ raise InvalidUrnError("DataProductUrn id cannot be empty")
158
+ if isinstance(id, DataProductUrn):
159
+ id = id.id
160
+ elif isinstance(id, Urn):
161
+ raise InvalidUrnError(f'Expecting a DataProductUrn but got {id}')
162
+ if UrnEncoder.contains_reserved_char(id):
163
+ raise InvalidUrnError(f'DataProductUrn id contains reserved characters')
131
164
 
132
- super().__init__(self.ENTITY_TYPE, [name])
165
+ super().__init__(self.ENTITY_TYPE, [id])
133
166
 
134
167
  @classmethod
135
- def _parse_ids(cls, entity_ids: List[str]) -> "CorpGroupUrn":
168
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataProductUrn":
136
169
  if len(entity_ids) != cls._URN_PARTS:
137
- raise InvalidUrnError(f"CorpGroupUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
138
- return cls(name=entity_ids[0], _allow_coercion=False)
170
+ raise InvalidUrnError(f"DataProductUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
171
+ return cls(id=entity_ids[0], _allow_coercion=False)
139
172
 
140
173
  @classmethod
141
- def underlying_key_aspect_type(cls) -> Type["CorpGroupKeyClass"]:
142
- from datahub.metadata.schema_classes import CorpGroupKeyClass
174
+ def underlying_key_aspect_type(cls) -> Type["DataProductKeyClass"]:
175
+ from datahub.metadata.schema_classes import DataProductKeyClass
143
176
 
144
- return CorpGroupKeyClass
177
+ return DataProductKeyClass
145
178
 
146
- def to_key_aspect(self) -> "CorpGroupKeyClass":
147
- from datahub.metadata.schema_classes import CorpGroupKeyClass
179
+ def to_key_aspect(self) -> "DataProductKeyClass":
180
+ from datahub.metadata.schema_classes import DataProductKeyClass
148
181
 
149
- return CorpGroupKeyClass(name=self.name)
182
+ return DataProductKeyClass(id=self.id)
150
183
 
151
184
  @classmethod
152
- def from_key_aspect(cls, key_aspect: "CorpGroupKeyClass") -> "CorpGroupUrn":
153
- return cls(name=key_aspect.name)
185
+ def from_key_aspect(cls, key_aspect: "DataProductKeyClass") -> "DataProductUrn":
186
+ return cls(id=key_aspect.id)
187
+
188
+ @property
189
+ def id(self) -> str:
190
+ return self.entity_ids[0]
191
+
192
+ if TYPE_CHECKING:
193
+ from datahub.metadata.schema_classes import DataTypeKeyClass
194
+
195
+ class DataTypeUrn(_SpecificUrn):
196
+ ENTITY_TYPE: ClassVar[str] = "dataType"
197
+ _URN_PARTS: ClassVar[int] = 1
198
+
199
+ def __init__(self, id: Union["DataTypeUrn", str], *, _allow_coercion: bool = True) -> None:
200
+ if _allow_coercion:
201
+ # Field coercion logic (if any is required).
202
+ if isinstance(id, str):
203
+ if id.startswith('urn:li:'):
204
+ try:
205
+ id = DataTypeUrn.from_string(id)
206
+ except InvalidUrnError:
207
+ raise InvalidUrnError(f'Expecting a DataTypeUrn but got {id}')
208
+ else:
209
+ id = UrnEncoder.encode_string(id)
210
+
211
+ # Validation logic.
212
+ if not id:
213
+ raise InvalidUrnError("DataTypeUrn id cannot be empty")
214
+ if isinstance(id, DataTypeUrn):
215
+ id = id.id
216
+ elif isinstance(id, Urn):
217
+ raise InvalidUrnError(f'Expecting a DataTypeUrn but got {id}')
218
+ if UrnEncoder.contains_reserved_char(id):
219
+ raise InvalidUrnError(f'DataTypeUrn id contains reserved characters')
220
+
221
+ super().__init__(self.ENTITY_TYPE, [id])
154
222
 
155
223
  @classmethod
156
- @deprecated(reason="Use the constructor instead")
157
- def create_from_id(cls, id: str) -> "CorpGroupUrn":
158
- return cls(id)
224
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataTypeUrn":
225
+ if len(entity_ids) != cls._URN_PARTS:
226
+ raise InvalidUrnError(f"DataTypeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
227
+ return cls(id=entity_ids[0], _allow_coercion=False)
228
+
229
+ @classmethod
230
+ def underlying_key_aspect_type(cls) -> Type["DataTypeKeyClass"]:
231
+ from datahub.metadata.schema_classes import DataTypeKeyClass
232
+
233
+ return DataTypeKeyClass
234
+
235
+ def to_key_aspect(self) -> "DataTypeKeyClass":
236
+ from datahub.metadata.schema_classes import DataTypeKeyClass
237
+
238
+ return DataTypeKeyClass(id=self.id)
239
+
240
+ @classmethod
241
+ def from_key_aspect(cls, key_aspect: "DataTypeKeyClass") -> "DataTypeUrn":
242
+ return cls(id=key_aspect.id)
159
243
 
160
244
  @property
161
- def name(self) -> str:
245
+ def id(self) -> str:
162
246
  return self.entity_ids[0]
163
247
 
164
248
  if TYPE_CHECKING:
165
- from datahub.metadata.schema_classes import GlobalSettingsKeyClass
249
+ from datahub.metadata.schema_classes import PlatformResourceKeyClass
166
250
 
167
- class GlobalSettingsUrn(_SpecificUrn):
168
- ENTITY_TYPE: ClassVar[str] = "globalSettings"
251
+ class PlatformResourceUrn(_SpecificUrn):
252
+ ENTITY_TYPE: ClassVar[str] = "platformResource"
169
253
  _URN_PARTS: ClassVar[int] = 1
170
254
 
171
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
255
+ def __init__(self, id: Union["PlatformResourceUrn", str], *, _allow_coercion: bool = True) -> None:
172
256
  if _allow_coercion:
173
257
  # Field coercion logic (if any is required).
174
- id = UrnEncoder.encode_string(id)
258
+ if isinstance(id, str):
259
+ if id.startswith('urn:li:'):
260
+ try:
261
+ id = PlatformResourceUrn.from_string(id)
262
+ except InvalidUrnError:
263
+ raise InvalidUrnError(f'Expecting a PlatformResourceUrn but got {id}')
264
+ else:
265
+ id = UrnEncoder.encode_string(id)
175
266
 
176
267
  # Validation logic.
177
268
  if not id:
178
- raise InvalidUrnError("GlobalSettingsUrn id cannot be empty")
269
+ raise InvalidUrnError("PlatformResourceUrn id cannot be empty")
270
+ if isinstance(id, PlatformResourceUrn):
271
+ id = id.id
272
+ elif isinstance(id, Urn):
273
+ raise InvalidUrnError(f'Expecting a PlatformResourceUrn but got {id}')
179
274
  if UrnEncoder.contains_reserved_char(id):
180
- raise InvalidUrnError(f'GlobalSettingsUrn id contains reserved characters')
275
+ raise InvalidUrnError(f'PlatformResourceUrn id contains reserved characters')
181
276
 
182
277
  super().__init__(self.ENTITY_TYPE, [id])
183
278
 
184
279
  @classmethod
185
- def _parse_ids(cls, entity_ids: List[str]) -> "GlobalSettingsUrn":
280
+ def _parse_ids(cls, entity_ids: List[str]) -> "PlatformResourceUrn":
186
281
  if len(entity_ids) != cls._URN_PARTS:
187
- raise InvalidUrnError(f"GlobalSettingsUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
282
+ raise InvalidUrnError(f"PlatformResourceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
188
283
  return cls(id=entity_ids[0], _allow_coercion=False)
189
284
 
190
285
  @classmethod
191
- def underlying_key_aspect_type(cls) -> Type["GlobalSettingsKeyClass"]:
192
- from datahub.metadata.schema_classes import GlobalSettingsKeyClass
286
+ def underlying_key_aspect_type(cls) -> Type["PlatformResourceKeyClass"]:
287
+ from datahub.metadata.schema_classes import PlatformResourceKeyClass
193
288
 
194
- return GlobalSettingsKeyClass
289
+ return PlatformResourceKeyClass
195
290
 
196
- def to_key_aspect(self) -> "GlobalSettingsKeyClass":
197
- from datahub.metadata.schema_classes import GlobalSettingsKeyClass
291
+ def to_key_aspect(self) -> "PlatformResourceKeyClass":
292
+ from datahub.metadata.schema_classes import PlatformResourceKeyClass
198
293
 
199
- return GlobalSettingsKeyClass(id=self.id)
294
+ return PlatformResourceKeyClass(id=self.id)
200
295
 
201
296
  @classmethod
202
- def from_key_aspect(cls, key_aspect: "GlobalSettingsKeyClass") -> "GlobalSettingsUrn":
297
+ def from_key_aspect(cls, key_aspect: "PlatformResourceKeyClass") -> "PlatformResourceUrn":
203
298
  return cls(id=key_aspect.id)
204
299
 
205
300
  @property
@@ -207,97 +302,118 @@ class GlobalSettingsUrn(_SpecificUrn):
207
302
  return self.entity_ids[0]
208
303
 
209
304
  if TYPE_CHECKING:
210
- from datahub.metadata.schema_classes import SchemaFieldKeyClass
305
+ from datahub.metadata.schema_classes import MLModelGroupKeyClass
211
306
 
212
- class SchemaFieldUrn(_SpecificUrn):
213
- ENTITY_TYPE: ClassVar[str] = "schemaField"
214
- _URN_PARTS: ClassVar[int] = 2
307
+ class MlModelGroupUrn(_SpecificUrn):
308
+ ENTITY_TYPE: ClassVar[str] = "mlModelGroup"
309
+ _URN_PARTS: ClassVar[int] = 3
215
310
 
216
- def __init__(self, parent: Union["Urn", str], field_path: str, *, _allow_coercion: bool = True) -> None:
311
+ def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
217
312
  if _allow_coercion:
218
313
  # Field coercion logic (if any is required).
219
- field_path = UrnEncoder.encode_string(field_path)
314
+ platform = DataPlatformUrn(platform).urn()
315
+ name = UrnEncoder.encode_string(name)
316
+ env = env.upper()
220
317
 
221
318
  # Validation logic.
222
- if not parent:
223
- raise InvalidUrnError("SchemaFieldUrn parent cannot be empty")
224
- parent = str(parent)
225
- assert Urn.from_string(parent)
226
- if not field_path:
227
- raise InvalidUrnError("SchemaFieldUrn field_path cannot be empty")
228
- if UrnEncoder.contains_reserved_char(field_path):
229
- raise InvalidUrnError(f'SchemaFieldUrn field_path contains reserved characters')
319
+ if not platform:
320
+ raise InvalidUrnError("MlModelGroupUrn platform cannot be empty")
321
+ platform = str(platform) # convert urn type to str
322
+ assert DataPlatformUrn.from_string(platform)
323
+ if not name:
324
+ raise InvalidUrnError("MlModelGroupUrn name cannot be empty")
325
+ if UrnEncoder.contains_reserved_char(name):
326
+ raise InvalidUrnError(f'MlModelGroupUrn name contains reserved characters')
327
+ if not env:
328
+ raise InvalidUrnError("MlModelGroupUrn env cannot be empty")
329
+ if UrnEncoder.contains_reserved_char(env):
330
+ raise InvalidUrnError(f'MlModelGroupUrn env contains reserved characters')
230
331
 
231
- super().__init__(self.ENTITY_TYPE, [parent, field_path])
332
+ super().__init__(self.ENTITY_TYPE, [platform, name, env])
232
333
 
233
334
  @classmethod
234
- def _parse_ids(cls, entity_ids: List[str]) -> "SchemaFieldUrn":
335
+ def _parse_ids(cls, entity_ids: List[str]) -> "MlModelGroupUrn":
235
336
  if len(entity_ids) != cls._URN_PARTS:
236
- raise InvalidUrnError(f"SchemaFieldUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
237
- return cls(parent=entity_ids[0], field_path=entity_ids[1], _allow_coercion=False)
337
+ raise InvalidUrnError(f"MlModelGroupUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
338
+ return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
238
339
 
239
340
  @classmethod
240
- def underlying_key_aspect_type(cls) -> Type["SchemaFieldKeyClass"]:
241
- from datahub.metadata.schema_classes import SchemaFieldKeyClass
341
+ def underlying_key_aspect_type(cls) -> Type["MLModelGroupKeyClass"]:
342
+ from datahub.metadata.schema_classes import MLModelGroupKeyClass
242
343
 
243
- return SchemaFieldKeyClass
344
+ return MLModelGroupKeyClass
244
345
 
245
- def to_key_aspect(self) -> "SchemaFieldKeyClass":
246
- from datahub.metadata.schema_classes import SchemaFieldKeyClass
346
+ def to_key_aspect(self) -> "MLModelGroupKeyClass":
347
+ from datahub.metadata.schema_classes import MLModelGroupKeyClass
247
348
 
248
- return SchemaFieldKeyClass(parent=self.parent, fieldPath=self.field_path)
349
+ return MLModelGroupKeyClass(platform=self.platform, name=self.name, origin=self.env)
249
350
 
250
351
  @classmethod
251
- def from_key_aspect(cls, key_aspect: "SchemaFieldKeyClass") -> "SchemaFieldUrn":
252
- return cls(parent=key_aspect.parent, field_path=key_aspect.fieldPath)
352
+ def from_key_aspect(cls, key_aspect: "MLModelGroupKeyClass") -> "MlModelGroupUrn":
353
+ return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
253
354
 
254
355
  @property
255
- def parent(self) -> str:
356
+ def platform(self) -> str:
256
357
  return self.entity_ids[0]
257
358
 
258
359
  @property
259
- def field_path(self) -> str:
360
+ def name(self) -> str:
260
361
  return self.entity_ids[1]
261
362
 
363
+ @property
364
+ def env(self) -> str:
365
+ return self.entity_ids[2]
366
+
262
367
  if TYPE_CHECKING:
263
- from datahub.metadata.schema_classes import DataHubStepStateKeyClass
368
+ from datahub.metadata.schema_classes import DataHubPolicyKeyClass
264
369
 
265
- class DataHubStepStateUrn(_SpecificUrn):
266
- ENTITY_TYPE: ClassVar[str] = "dataHubStepState"
370
+ class DataHubPolicyUrn(_SpecificUrn):
371
+ ENTITY_TYPE: ClassVar[str] = "dataHubPolicy"
267
372
  _URN_PARTS: ClassVar[int] = 1
268
373
 
269
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
374
+ def __init__(self, id: Union["DataHubPolicyUrn", str], *, _allow_coercion: bool = True) -> None:
270
375
  if _allow_coercion:
271
376
  # Field coercion logic (if any is required).
272
- id = UrnEncoder.encode_string(id)
377
+ if isinstance(id, str):
378
+ if id.startswith('urn:li:'):
379
+ try:
380
+ id = DataHubPolicyUrn.from_string(id)
381
+ except InvalidUrnError:
382
+ raise InvalidUrnError(f'Expecting a DataHubPolicyUrn but got {id}')
383
+ else:
384
+ id = UrnEncoder.encode_string(id)
273
385
 
274
386
  # Validation logic.
275
387
  if not id:
276
- raise InvalidUrnError("DataHubStepStateUrn id cannot be empty")
388
+ raise InvalidUrnError("DataHubPolicyUrn id cannot be empty")
389
+ if isinstance(id, DataHubPolicyUrn):
390
+ id = id.id
391
+ elif isinstance(id, Urn):
392
+ raise InvalidUrnError(f'Expecting a DataHubPolicyUrn but got {id}')
277
393
  if UrnEncoder.contains_reserved_char(id):
278
- raise InvalidUrnError(f'DataHubStepStateUrn id contains reserved characters')
394
+ raise InvalidUrnError(f'DataHubPolicyUrn id contains reserved characters')
279
395
 
280
396
  super().__init__(self.ENTITY_TYPE, [id])
281
397
 
282
398
  @classmethod
283
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubStepStateUrn":
399
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubPolicyUrn":
284
400
  if len(entity_ids) != cls._URN_PARTS:
285
- raise InvalidUrnError(f"DataHubStepStateUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
401
+ raise InvalidUrnError(f"DataHubPolicyUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
286
402
  return cls(id=entity_ids[0], _allow_coercion=False)
287
403
 
288
404
  @classmethod
289
- def underlying_key_aspect_type(cls) -> Type["DataHubStepStateKeyClass"]:
290
- from datahub.metadata.schema_classes import DataHubStepStateKeyClass
405
+ def underlying_key_aspect_type(cls) -> Type["DataHubPolicyKeyClass"]:
406
+ from datahub.metadata.schema_classes import DataHubPolicyKeyClass
291
407
 
292
- return DataHubStepStateKeyClass
408
+ return DataHubPolicyKeyClass
293
409
 
294
- def to_key_aspect(self) -> "DataHubStepStateKeyClass":
295
- from datahub.metadata.schema_classes import DataHubStepStateKeyClass
410
+ def to_key_aspect(self) -> "DataHubPolicyKeyClass":
411
+ from datahub.metadata.schema_classes import DataHubPolicyKeyClass
296
412
 
297
- return DataHubStepStateKeyClass(id=self.id)
413
+ return DataHubPolicyKeyClass(id=self.id)
298
414
 
299
415
  @classmethod
300
- def from_key_aspect(cls, key_aspect: "DataHubStepStateKeyClass") -> "DataHubStepStateUrn":
416
+ def from_key_aspect(cls, key_aspect: "DataHubPolicyKeyClass") -> "DataHubPolicyUrn":
301
417
  return cls(id=key_aspect.id)
302
418
 
303
419
  @property
@@ -305,154 +421,185 @@ class DataHubStepStateUrn(_SpecificUrn):
305
421
  return self.entity_ids[0]
306
422
 
307
423
  if TYPE_CHECKING:
308
- from datahub.metadata.schema_classes import CorpUserKeyClass
424
+ from datahub.metadata.schema_classes import TagKeyClass
309
425
 
310
- class CorpUserUrn(_SpecificUrn):
311
- ENTITY_TYPE: ClassVar[str] = "corpuser"
426
+ class TagUrn(_SpecificUrn):
427
+ ENTITY_TYPE: ClassVar[str] = "tag"
312
428
  _URN_PARTS: ClassVar[int] = 1
313
429
 
314
- def __init__(self, username: str, *, _allow_coercion: bool = True) -> None:
430
+ def __init__(self, name: Union["TagUrn", str], *, _allow_coercion: bool = True) -> None:
315
431
  if _allow_coercion:
316
432
  # Field coercion logic (if any is required).
317
- username = UrnEncoder.encode_string(username)
433
+ if isinstance(name, str):
434
+ if name.startswith('urn:li:'):
435
+ try:
436
+ name = TagUrn.from_string(name)
437
+ except InvalidUrnError:
438
+ raise InvalidUrnError(f'Expecting a TagUrn but got {name}')
439
+ else:
440
+ name = UrnEncoder.encode_string(name)
318
441
 
319
442
  # Validation logic.
320
- if not username:
321
- raise InvalidUrnError("CorpUserUrn username cannot be empty")
322
- if UrnEncoder.contains_reserved_char(username):
323
- raise InvalidUrnError(f'CorpUserUrn username contains reserved characters')
443
+ if not name:
444
+ raise InvalidUrnError("TagUrn name cannot be empty")
445
+ if isinstance(name, TagUrn):
446
+ name = name.name
447
+ elif isinstance(name, Urn):
448
+ raise InvalidUrnError(f'Expecting a TagUrn but got {name}')
449
+ if UrnEncoder.contains_reserved_char(name):
450
+ raise InvalidUrnError(f'TagUrn name contains reserved characters')
324
451
 
325
- super().__init__(self.ENTITY_TYPE, [username])
452
+ super().__init__(self.ENTITY_TYPE, [name])
326
453
 
327
454
  @classmethod
328
- def _parse_ids(cls, entity_ids: List[str]) -> "CorpUserUrn":
455
+ def _parse_ids(cls, entity_ids: List[str]) -> "TagUrn":
329
456
  if len(entity_ids) != cls._URN_PARTS:
330
- raise InvalidUrnError(f"CorpUserUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
331
- return cls(username=entity_ids[0], _allow_coercion=False)
457
+ raise InvalidUrnError(f"TagUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
458
+ return cls(name=entity_ids[0], _allow_coercion=False)
332
459
 
333
460
  @classmethod
334
- def underlying_key_aspect_type(cls) -> Type["CorpUserKeyClass"]:
335
- from datahub.metadata.schema_classes import CorpUserKeyClass
461
+ def underlying_key_aspect_type(cls) -> Type["TagKeyClass"]:
462
+ from datahub.metadata.schema_classes import TagKeyClass
336
463
 
337
- return CorpUserKeyClass
464
+ return TagKeyClass
338
465
 
339
- def to_key_aspect(self) -> "CorpUserKeyClass":
340
- from datahub.metadata.schema_classes import CorpUserKeyClass
466
+ def to_key_aspect(self) -> "TagKeyClass":
467
+ from datahub.metadata.schema_classes import TagKeyClass
341
468
 
342
- return CorpUserKeyClass(username=self.username)
469
+ return TagKeyClass(name=self.name)
343
470
 
344
471
  @classmethod
345
- def from_key_aspect(cls, key_aspect: "CorpUserKeyClass") -> "CorpUserUrn":
346
- return cls(username=key_aspect.username)
472
+ def from_key_aspect(cls, key_aspect: "TagKeyClass") -> "TagUrn":
473
+ return cls(name=key_aspect.name)
347
474
 
348
475
  @classmethod
349
476
  @deprecated(reason="Use the constructor instead")
350
- def create_from_id(cls, id: str) -> "CorpUserUrn":
477
+ def create_from_id(cls, id: str) -> "TagUrn":
351
478
  return cls(id)
352
479
 
353
480
  @property
354
- def username(self) -> str:
481
+ def name(self) -> str:
355
482
  return self.entity_ids[0]
356
483
 
357
484
  if TYPE_CHECKING:
358
- from datahub.metadata.schema_classes import DataHubConnectionKeyClass
485
+ from datahub.metadata.schema_classes import MLModelDeploymentKeyClass
359
486
 
360
- class DataHubConnectionUrn(_SpecificUrn):
361
- ENTITY_TYPE: ClassVar[str] = "dataHubConnection"
362
- _URN_PARTS: ClassVar[int] = 1
487
+ class MlModelDeploymentUrn(_SpecificUrn):
488
+ ENTITY_TYPE: ClassVar[str] = "mlModelDeployment"
489
+ _URN_PARTS: ClassVar[int] = 3
363
490
 
364
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
491
+ def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
365
492
  if _allow_coercion:
366
493
  # Field coercion logic (if any is required).
367
- id = UrnEncoder.encode_string(id)
494
+ platform = DataPlatformUrn(platform).urn()
495
+ name = UrnEncoder.encode_string(name)
496
+ env = env.upper()
368
497
 
369
498
  # Validation logic.
370
- if not id:
371
- raise InvalidUrnError("DataHubConnectionUrn id cannot be empty")
372
- if UrnEncoder.contains_reserved_char(id):
373
- raise InvalidUrnError(f'DataHubConnectionUrn id contains reserved characters')
499
+ if not platform:
500
+ raise InvalidUrnError("MlModelDeploymentUrn platform cannot be empty")
501
+ platform = str(platform) # convert urn type to str
502
+ assert DataPlatformUrn.from_string(platform)
503
+ if not name:
504
+ raise InvalidUrnError("MlModelDeploymentUrn name cannot be empty")
505
+ if UrnEncoder.contains_reserved_char(name):
506
+ raise InvalidUrnError(f'MlModelDeploymentUrn name contains reserved characters')
507
+ if not env:
508
+ raise InvalidUrnError("MlModelDeploymentUrn env cannot be empty")
509
+ if UrnEncoder.contains_reserved_char(env):
510
+ raise InvalidUrnError(f'MlModelDeploymentUrn env contains reserved characters')
374
511
 
375
- super().__init__(self.ENTITY_TYPE, [id])
512
+ super().__init__(self.ENTITY_TYPE, [platform, name, env])
376
513
 
377
514
  @classmethod
378
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubConnectionUrn":
515
+ def _parse_ids(cls, entity_ids: List[str]) -> "MlModelDeploymentUrn":
379
516
  if len(entity_ids) != cls._URN_PARTS:
380
- raise InvalidUrnError(f"DataHubConnectionUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
381
- return cls(id=entity_ids[0], _allow_coercion=False)
517
+ raise InvalidUrnError(f"MlModelDeploymentUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
518
+ return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
382
519
 
383
520
  @classmethod
384
- def underlying_key_aspect_type(cls) -> Type["DataHubConnectionKeyClass"]:
385
- from datahub.metadata.schema_classes import DataHubConnectionKeyClass
521
+ def underlying_key_aspect_type(cls) -> Type["MLModelDeploymentKeyClass"]:
522
+ from datahub.metadata.schema_classes import MLModelDeploymentKeyClass
386
523
 
387
- return DataHubConnectionKeyClass
524
+ return MLModelDeploymentKeyClass
388
525
 
389
- def to_key_aspect(self) -> "DataHubConnectionKeyClass":
390
- from datahub.metadata.schema_classes import DataHubConnectionKeyClass
526
+ def to_key_aspect(self) -> "MLModelDeploymentKeyClass":
527
+ from datahub.metadata.schema_classes import MLModelDeploymentKeyClass
391
528
 
392
- return DataHubConnectionKeyClass(id=self.id)
529
+ return MLModelDeploymentKeyClass(platform=self.platform, name=self.name, origin=self.env)
393
530
 
394
531
  @classmethod
395
- def from_key_aspect(cls, key_aspect: "DataHubConnectionKeyClass") -> "DataHubConnectionUrn":
396
- return cls(id=key_aspect.id)
532
+ def from_key_aspect(cls, key_aspect: "MLModelDeploymentKeyClass") -> "MlModelDeploymentUrn":
533
+ return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
397
534
 
398
535
  @property
399
- def id(self) -> str:
536
+ def platform(self) -> str:
400
537
  return self.entity_ids[0]
401
538
 
539
+ @property
540
+ def name(self) -> str:
541
+ return self.entity_ids[1]
542
+
543
+ @property
544
+ def env(self) -> str:
545
+ return self.entity_ids[2]
546
+
402
547
  if TYPE_CHECKING:
403
- from datahub.metadata.schema_classes import DashboardKeyClass
548
+ from datahub.metadata.schema_classes import DataHubActionKeyClass
404
549
 
405
- class DashboardUrn(_SpecificUrn):
406
- ENTITY_TYPE: ClassVar[str] = "dashboard"
407
- _URN_PARTS: ClassVar[int] = 2
550
+ class DataHubActionUrn(_SpecificUrn):
551
+ ENTITY_TYPE: ClassVar[str] = "dataHubAction"
552
+ _URN_PARTS: ClassVar[int] = 1
408
553
 
409
- def __init__(self, dashboard_tool: str, dashboard_id: str, *, _allow_coercion: bool = True) -> None:
554
+ def __init__(self, id: Union["DataHubActionUrn", str], *, _allow_coercion: bool = True) -> None:
410
555
  if _allow_coercion:
411
556
  # Field coercion logic (if any is required).
412
- dashboard_tool = UrnEncoder.encode_string(dashboard_tool)
413
- dashboard_id = UrnEncoder.encode_string(dashboard_id)
557
+ if isinstance(id, str):
558
+ if id.startswith('urn:li:'):
559
+ try:
560
+ id = DataHubActionUrn.from_string(id)
561
+ except InvalidUrnError:
562
+ raise InvalidUrnError(f'Expecting a DataHubActionUrn but got {id}')
563
+ else:
564
+ id = UrnEncoder.encode_string(id)
414
565
 
415
566
  # Validation logic.
416
- if not dashboard_tool:
417
- raise InvalidUrnError("DashboardUrn dashboard_tool cannot be empty")
418
- if UrnEncoder.contains_reserved_char(dashboard_tool):
419
- raise InvalidUrnError(f'DashboardUrn dashboard_tool contains reserved characters')
420
- if not dashboard_id:
421
- raise InvalidUrnError("DashboardUrn dashboard_id cannot be empty")
422
- if UrnEncoder.contains_reserved_char(dashboard_id):
423
- raise InvalidUrnError(f'DashboardUrn dashboard_id contains reserved characters')
567
+ if not id:
568
+ raise InvalidUrnError("DataHubActionUrn id cannot be empty")
569
+ if isinstance(id, DataHubActionUrn):
570
+ id = id.id
571
+ elif isinstance(id, Urn):
572
+ raise InvalidUrnError(f'Expecting a DataHubActionUrn but got {id}')
573
+ if UrnEncoder.contains_reserved_char(id):
574
+ raise InvalidUrnError(f'DataHubActionUrn id contains reserved characters')
424
575
 
425
- super().__init__(self.ENTITY_TYPE, [dashboard_tool, dashboard_id])
576
+ super().__init__(self.ENTITY_TYPE, [id])
426
577
 
427
578
  @classmethod
428
- def _parse_ids(cls, entity_ids: List[str]) -> "DashboardUrn":
579
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubActionUrn":
429
580
  if len(entity_ids) != cls._URN_PARTS:
430
- raise InvalidUrnError(f"DashboardUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
431
- return cls(dashboard_tool=entity_ids[0], dashboard_id=entity_ids[1], _allow_coercion=False)
581
+ raise InvalidUrnError(f"DataHubActionUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
582
+ return cls(id=entity_ids[0], _allow_coercion=False)
432
583
 
433
584
  @classmethod
434
- def underlying_key_aspect_type(cls) -> Type["DashboardKeyClass"]:
435
- from datahub.metadata.schema_classes import DashboardKeyClass
585
+ def underlying_key_aspect_type(cls) -> Type["DataHubActionKeyClass"]:
586
+ from datahub.metadata.schema_classes import DataHubActionKeyClass
436
587
 
437
- return DashboardKeyClass
588
+ return DataHubActionKeyClass
438
589
 
439
- def to_key_aspect(self) -> "DashboardKeyClass":
440
- from datahub.metadata.schema_classes import DashboardKeyClass
590
+ def to_key_aspect(self) -> "DataHubActionKeyClass":
591
+ from datahub.metadata.schema_classes import DataHubActionKeyClass
441
592
 
442
- return DashboardKeyClass(dashboardTool=self.dashboard_tool, dashboardId=self.dashboard_id)
593
+ return DataHubActionKeyClass(id=self.id)
443
594
 
444
595
  @classmethod
445
- def from_key_aspect(cls, key_aspect: "DashboardKeyClass") -> "DashboardUrn":
446
- return cls(dashboard_tool=key_aspect.dashboardTool, dashboard_id=key_aspect.dashboardId)
596
+ def from_key_aspect(cls, key_aspect: "DataHubActionKeyClass") -> "DataHubActionUrn":
597
+ return cls(id=key_aspect.id)
447
598
 
448
599
  @property
449
- def dashboard_tool(self) -> str:
600
+ def id(self) -> str:
450
601
  return self.entity_ids[0]
451
602
 
452
- @property
453
- def dashboard_id(self) -> str:
454
- return self.entity_ids[1]
455
-
456
603
  if TYPE_CHECKING:
457
604
  from datahub.metadata.schema_classes import DataHubRoleKeyClass
458
605
 
@@ -460,14 +607,25 @@ class DataHubRoleUrn(_SpecificUrn):
460
607
  ENTITY_TYPE: ClassVar[str] = "dataHubRole"
461
608
  _URN_PARTS: ClassVar[int] = 1
462
609
 
463
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
610
+ def __init__(self, id: Union["DataHubRoleUrn", str], *, _allow_coercion: bool = True) -> None:
464
611
  if _allow_coercion:
465
612
  # Field coercion logic (if any is required).
466
- id = UrnEncoder.encode_string(id)
613
+ if isinstance(id, str):
614
+ if id.startswith('urn:li:'):
615
+ try:
616
+ id = DataHubRoleUrn.from_string(id)
617
+ except InvalidUrnError:
618
+ raise InvalidUrnError(f'Expecting a DataHubRoleUrn but got {id}')
619
+ else:
620
+ id = UrnEncoder.encode_string(id)
467
621
 
468
622
  # Validation logic.
469
623
  if not id:
470
624
  raise InvalidUrnError("DataHubRoleUrn id cannot be empty")
625
+ if isinstance(id, DataHubRoleUrn):
626
+ id = id.id
627
+ elif isinstance(id, Urn):
628
+ raise InvalidUrnError(f'Expecting a DataHubRoleUrn but got {id}')
471
629
  if UrnEncoder.contains_reserved_char(id):
472
630
  raise InvalidUrnError(f'DataHubRoleUrn id contains reserved characters')
473
631
 
@@ -499,44 +657,55 @@ class DataHubRoleUrn(_SpecificUrn):
499
657
  return self.entity_ids[0]
500
658
 
501
659
  if TYPE_CHECKING:
502
- from datahub.metadata.schema_classes import ExecutionRequestKeyClass
660
+ from datahub.metadata.schema_classes import DataHubUpgradeKeyClass
503
661
 
504
- class DataHubExecutionRequestUrn(_SpecificUrn):
505
- ENTITY_TYPE: ClassVar[str] = "dataHubExecutionRequest"
662
+ class DataHubUpgradeUrn(_SpecificUrn):
663
+ ENTITY_TYPE: ClassVar[str] = "dataHubUpgrade"
506
664
  _URN_PARTS: ClassVar[int] = 1
507
665
 
508
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
666
+ def __init__(self, id: Union["DataHubUpgradeUrn", str], *, _allow_coercion: bool = True) -> None:
509
667
  if _allow_coercion:
510
668
  # Field coercion logic (if any is required).
511
- id = UrnEncoder.encode_string(id)
669
+ if isinstance(id, str):
670
+ if id.startswith('urn:li:'):
671
+ try:
672
+ id = DataHubUpgradeUrn.from_string(id)
673
+ except InvalidUrnError:
674
+ raise InvalidUrnError(f'Expecting a DataHubUpgradeUrn but got {id}')
675
+ else:
676
+ id = UrnEncoder.encode_string(id)
512
677
 
513
678
  # Validation logic.
514
679
  if not id:
515
- raise InvalidUrnError("DataHubExecutionRequestUrn id cannot be empty")
680
+ raise InvalidUrnError("DataHubUpgradeUrn id cannot be empty")
681
+ if isinstance(id, DataHubUpgradeUrn):
682
+ id = id.id
683
+ elif isinstance(id, Urn):
684
+ raise InvalidUrnError(f'Expecting a DataHubUpgradeUrn but got {id}')
516
685
  if UrnEncoder.contains_reserved_char(id):
517
- raise InvalidUrnError(f'DataHubExecutionRequestUrn id contains reserved characters')
686
+ raise InvalidUrnError(f'DataHubUpgradeUrn id contains reserved characters')
518
687
 
519
688
  super().__init__(self.ENTITY_TYPE, [id])
520
689
 
521
690
  @classmethod
522
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubExecutionRequestUrn":
691
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubUpgradeUrn":
523
692
  if len(entity_ids) != cls._URN_PARTS:
524
- raise InvalidUrnError(f"DataHubExecutionRequestUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
693
+ raise InvalidUrnError(f"DataHubUpgradeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
525
694
  return cls(id=entity_ids[0], _allow_coercion=False)
526
695
 
527
696
  @classmethod
528
- def underlying_key_aspect_type(cls) -> Type["ExecutionRequestKeyClass"]:
529
- from datahub.metadata.schema_classes import ExecutionRequestKeyClass
697
+ def underlying_key_aspect_type(cls) -> Type["DataHubUpgradeKeyClass"]:
698
+ from datahub.metadata.schema_classes import DataHubUpgradeKeyClass
530
699
 
531
- return ExecutionRequestKeyClass
700
+ return DataHubUpgradeKeyClass
532
701
 
533
- def to_key_aspect(self) -> "ExecutionRequestKeyClass":
534
- from datahub.metadata.schema_classes import ExecutionRequestKeyClass
702
+ def to_key_aspect(self) -> "DataHubUpgradeKeyClass":
703
+ from datahub.metadata.schema_classes import DataHubUpgradeKeyClass
535
704
 
536
- return ExecutionRequestKeyClass(id=self.id)
705
+ return DataHubUpgradeKeyClass(id=self.id)
537
706
 
538
707
  @classmethod
539
- def from_key_aspect(cls, key_aspect: "ExecutionRequestKeyClass") -> "DataHubExecutionRequestUrn":
708
+ def from_key_aspect(cls, key_aspect: "DataHubUpgradeKeyClass") -> "DataHubUpgradeUrn":
540
709
  return cls(id=key_aspect.id)
541
710
 
542
711
  @property
@@ -544,415 +713,229 @@ class DataHubExecutionRequestUrn(_SpecificUrn):
544
713
  return self.entity_ids[0]
545
714
 
546
715
  if TYPE_CHECKING:
547
- from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
716
+ from datahub.metadata.schema_classes import FormKeyClass
548
717
 
549
- class DataProcessInstanceUrn(_SpecificUrn):
550
- ENTITY_TYPE: ClassVar[str] = "dataProcessInstance"
718
+ class FormUrn(_SpecificUrn):
719
+ ENTITY_TYPE: ClassVar[str] = "form"
551
720
  _URN_PARTS: ClassVar[int] = 1
552
721
 
553
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
722
+ def __init__(self, id: Union["FormUrn", str], *, _allow_coercion: bool = True) -> None:
554
723
  if _allow_coercion:
555
724
  # Field coercion logic (if any is required).
556
- id = UrnEncoder.encode_string(id)
725
+ if isinstance(id, str):
726
+ if id.startswith('urn:li:'):
727
+ try:
728
+ id = FormUrn.from_string(id)
729
+ except InvalidUrnError:
730
+ raise InvalidUrnError(f'Expecting a FormUrn but got {id}')
731
+ else:
732
+ id = UrnEncoder.encode_string(id)
557
733
 
558
734
  # Validation logic.
559
735
  if not id:
560
- raise InvalidUrnError("DataProcessInstanceUrn id cannot be empty")
736
+ raise InvalidUrnError("FormUrn id cannot be empty")
737
+ if isinstance(id, FormUrn):
738
+ id = id.id
739
+ elif isinstance(id, Urn):
740
+ raise InvalidUrnError(f'Expecting a FormUrn but got {id}')
561
741
  if UrnEncoder.contains_reserved_char(id):
562
- raise InvalidUrnError(f'DataProcessInstanceUrn id contains reserved characters')
742
+ raise InvalidUrnError(f'FormUrn id contains reserved characters')
563
743
 
564
744
  super().__init__(self.ENTITY_TYPE, [id])
565
745
 
566
746
  @classmethod
567
- def _parse_ids(cls, entity_ids: List[str]) -> "DataProcessInstanceUrn":
747
+ def _parse_ids(cls, entity_ids: List[str]) -> "FormUrn":
568
748
  if len(entity_ids) != cls._URN_PARTS:
569
- raise InvalidUrnError(f"DataProcessInstanceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
749
+ raise InvalidUrnError(f"FormUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
570
750
  return cls(id=entity_ids[0], _allow_coercion=False)
571
751
 
572
752
  @classmethod
573
- def underlying_key_aspect_type(cls) -> Type["DataProcessInstanceKeyClass"]:
574
- from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
753
+ def underlying_key_aspect_type(cls) -> Type["FormKeyClass"]:
754
+ from datahub.metadata.schema_classes import FormKeyClass
575
755
 
576
- return DataProcessInstanceKeyClass
756
+ return FormKeyClass
577
757
 
578
- def to_key_aspect(self) -> "DataProcessInstanceKeyClass":
579
- from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
758
+ def to_key_aspect(self) -> "FormKeyClass":
759
+ from datahub.metadata.schema_classes import FormKeyClass
580
760
 
581
- return DataProcessInstanceKeyClass(id=self.id)
761
+ return FormKeyClass(id=self.id)
582
762
 
583
763
  @classmethod
584
- def from_key_aspect(cls, key_aspect: "DataProcessInstanceKeyClass") -> "DataProcessInstanceUrn":
764
+ def from_key_aspect(cls, key_aspect: "FormKeyClass") -> "FormUrn":
585
765
  return cls(id=key_aspect.id)
586
766
 
587
- @classmethod
588
- @deprecated(reason="Use the constructor instead")
589
- def create_from_id(cls, id: str) -> "DataProcessInstanceUrn":
590
- return cls(id)
591
-
592
- @deprecated(reason="Use .id instead")
593
- def get_dataprocessinstance_id(self) -> str:
594
- return self.id
595
-
596
767
  @property
597
768
  def id(self) -> str:
598
769
  return self.entity_ids[0]
599
770
 
600
771
  if TYPE_CHECKING:
601
- from datahub.metadata.schema_classes import DataJobKeyClass
772
+ from datahub.metadata.schema_classes import NotebookKeyClass
602
773
 
603
- class DataJobUrn(_SpecificUrn):
604
- ENTITY_TYPE: ClassVar[str] = "dataJob"
774
+ class NotebookUrn(_SpecificUrn):
775
+ ENTITY_TYPE: ClassVar[str] = "notebook"
605
776
  _URN_PARTS: ClassVar[int] = 2
606
777
 
607
- def __init__(self, flow: Union["DataFlowUrn", str], job_id: str, *, _allow_coercion: bool = True) -> None:
778
+ def __init__(self, notebook_tool: str, notebook_id: str, *, _allow_coercion: bool = True) -> None:
608
779
  if _allow_coercion:
609
780
  # Field coercion logic (if any is required).
610
- job_id = UrnEncoder.encode_string(job_id)
781
+ notebook_tool = UrnEncoder.encode_string(notebook_tool)
782
+ notebook_id = UrnEncoder.encode_string(notebook_id)
611
783
 
612
784
  # Validation logic.
613
- if not flow:
614
- raise InvalidUrnError("DataJobUrn flow cannot be empty")
615
- flow = str(flow)
616
- assert DataFlowUrn.from_string(flow)
617
- if not job_id:
618
- raise InvalidUrnError("DataJobUrn job_id cannot be empty")
619
- if UrnEncoder.contains_reserved_char(job_id):
620
- raise InvalidUrnError(f'DataJobUrn job_id contains reserved characters')
785
+ if not notebook_tool:
786
+ raise InvalidUrnError("NotebookUrn notebook_tool cannot be empty")
787
+ if UrnEncoder.contains_reserved_char(notebook_tool):
788
+ raise InvalidUrnError(f'NotebookUrn notebook_tool contains reserved characters')
789
+ if not notebook_id:
790
+ raise InvalidUrnError("NotebookUrn notebook_id cannot be empty")
791
+ if UrnEncoder.contains_reserved_char(notebook_id):
792
+ raise InvalidUrnError(f'NotebookUrn notebook_id contains reserved characters')
621
793
 
622
- super().__init__(self.ENTITY_TYPE, [flow, job_id])
794
+ super().__init__(self.ENTITY_TYPE, [notebook_tool, notebook_id])
623
795
 
624
796
  @classmethod
625
- def _parse_ids(cls, entity_ids: List[str]) -> "DataJobUrn":
797
+ def _parse_ids(cls, entity_ids: List[str]) -> "NotebookUrn":
626
798
  if len(entity_ids) != cls._URN_PARTS:
627
- raise InvalidUrnError(f"DataJobUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
628
- return cls(flow=entity_ids[0], job_id=entity_ids[1], _allow_coercion=False)
799
+ raise InvalidUrnError(f"NotebookUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
800
+ return cls(notebook_tool=entity_ids[0], notebook_id=entity_ids[1], _allow_coercion=False)
629
801
 
630
802
  @classmethod
631
- def underlying_key_aspect_type(cls) -> Type["DataJobKeyClass"]:
632
- from datahub.metadata.schema_classes import DataJobKeyClass
803
+ def underlying_key_aspect_type(cls) -> Type["NotebookKeyClass"]:
804
+ from datahub.metadata.schema_classes import NotebookKeyClass
633
805
 
634
- return DataJobKeyClass
806
+ return NotebookKeyClass
635
807
 
636
- def to_key_aspect(self) -> "DataJobKeyClass":
637
- from datahub.metadata.schema_classes import DataJobKeyClass
808
+ def to_key_aspect(self) -> "NotebookKeyClass":
809
+ from datahub.metadata.schema_classes import NotebookKeyClass
638
810
 
639
- return DataJobKeyClass(flow=self.flow, jobId=self.job_id)
811
+ return NotebookKeyClass(notebookTool=self.notebook_tool, notebookId=self.notebook_id)
640
812
 
641
813
  @classmethod
642
- def from_key_aspect(cls, key_aspect: "DataJobKeyClass") -> "DataJobUrn":
643
- return cls(flow=key_aspect.flow, job_id=key_aspect.jobId)
814
+ def from_key_aspect(cls, key_aspect: "NotebookKeyClass") -> "NotebookUrn":
815
+ return cls(notebook_tool=key_aspect.notebookTool, notebook_id=key_aspect.notebookId)
644
816
 
645
- @classmethod
646
- def create_from_ids(cls, data_flow_urn: str, job_id: str) -> "DataJobUrn":
647
- return cls(data_flow_urn, job_id)
648
-
649
- def get_data_flow_urn(self) -> "DataFlowUrn":
650
- return DataFlowUrn.from_string(self.flow)
651
-
652
- @deprecated(reason="Use .job_id instead")
653
- def get_job_id(self) -> str:
654
- return self.job_id
655
-
656
- @property
657
- def flow(self) -> str:
658
- return self.entity_ids[0]
659
-
660
- @property
661
- def job_id(self) -> str:
662
- return self.entity_ids[1]
663
-
664
- if TYPE_CHECKING:
665
- from datahub.metadata.schema_classes import DataFlowKeyClass
666
-
667
- class DataFlowUrn(_SpecificUrn):
668
- ENTITY_TYPE: ClassVar[str] = "dataFlow"
669
- _URN_PARTS: ClassVar[int] = 3
670
-
671
- def __init__(self, orchestrator: str, flow_id: str, cluster: str, *, _allow_coercion: bool = True) -> None:
672
- if _allow_coercion:
673
- # Field coercion logic (if any is required).
674
- orchestrator = UrnEncoder.encode_string(orchestrator)
675
- flow_id = UrnEncoder.encode_string(flow_id)
676
- cluster = UrnEncoder.encode_string(cluster)
677
-
678
- # Validation logic.
679
- if not orchestrator:
680
- raise InvalidUrnError("DataFlowUrn orchestrator cannot be empty")
681
- if UrnEncoder.contains_reserved_char(orchestrator):
682
- raise InvalidUrnError(f'DataFlowUrn orchestrator contains reserved characters')
683
- if not flow_id:
684
- raise InvalidUrnError("DataFlowUrn flow_id cannot be empty")
685
- if UrnEncoder.contains_reserved_char(flow_id):
686
- raise InvalidUrnError(f'DataFlowUrn flow_id contains reserved characters')
687
- if not cluster:
688
- raise InvalidUrnError("DataFlowUrn cluster cannot be empty")
689
- if UrnEncoder.contains_reserved_char(cluster):
690
- raise InvalidUrnError(f'DataFlowUrn cluster contains reserved characters')
691
-
692
- super().__init__(self.ENTITY_TYPE, [orchestrator, flow_id, cluster])
693
-
694
- @classmethod
695
- def _parse_ids(cls, entity_ids: List[str]) -> "DataFlowUrn":
696
- if len(entity_ids) != cls._URN_PARTS:
697
- raise InvalidUrnError(f"DataFlowUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
698
- return cls(orchestrator=entity_ids[0], flow_id=entity_ids[1], cluster=entity_ids[2], _allow_coercion=False)
699
-
700
- @classmethod
701
- def underlying_key_aspect_type(cls) -> Type["DataFlowKeyClass"]:
702
- from datahub.metadata.schema_classes import DataFlowKeyClass
703
-
704
- return DataFlowKeyClass
705
-
706
- def to_key_aspect(self) -> "DataFlowKeyClass":
707
- from datahub.metadata.schema_classes import DataFlowKeyClass
708
-
709
- return DataFlowKeyClass(orchestrator=self.orchestrator, flowId=self.flow_id, cluster=self.cluster)
710
-
711
- @classmethod
712
- def from_key_aspect(cls, key_aspect: "DataFlowKeyClass") -> "DataFlowUrn":
713
- return cls(orchestrator=key_aspect.orchestrator, flow_id=key_aspect.flowId, cluster=key_aspect.cluster)
714
-
715
- @classmethod
716
- def create_from_ids(
717
- cls,
718
- orchestrator: str,
719
- flow_id: str,
720
- env: str,
721
- platform_instance: Optional[str] = None,
722
- ) -> "DataFlowUrn":
723
- return cls(
724
- orchestrator=orchestrator,
725
- flow_id=f"{platform_instance}.{flow_id}" if platform_instance else flow_id,
726
- cluster=env,
727
- )
728
-
729
- @deprecated(reason="Use .orchestrator instead")
730
- def get_orchestrator_name(self) -> str:
731
- return self.orchestrator
732
-
733
- @deprecated(reason="Use .flow_id instead")
734
- def get_flow_id(self) -> str:
735
- return self.flow_id
817
+ @deprecated(reason="Use .notebook_tool instead")
818
+ def get_platform_id(self) -> str:
819
+ return self.notebook_tool
736
820
 
737
- @deprecated(reason="Use .cluster instead")
738
- def get_env(self) -> str:
739
- return self.cluster
821
+ @deprecated(reason="Use .notebook_id instead")
822
+ def get_notebook_id(self) -> str:
823
+ return self.notebook_id
740
824
 
741
825
  @property
742
- def orchestrator(self) -> str:
826
+ def notebook_tool(self) -> str:
743
827
  return self.entity_ids[0]
744
828
 
745
829
  @property
746
- def flow_id(self) -> str:
830
+ def notebook_id(self) -> str:
747
831
  return self.entity_ids[1]
748
832
 
749
- @property
750
- def cluster(self) -> str:
751
- return self.entity_ids[2]
752
-
753
- if TYPE_CHECKING:
754
- from datahub.metadata.schema_classes import DataPlatformKeyClass
755
-
756
- class DataPlatformUrn(_SpecificUrn):
757
- ENTITY_TYPE: ClassVar[str] = "dataPlatform"
758
- _URN_PARTS: ClassVar[int] = 1
759
-
760
- def __init__(self, platform_name: str, *, _allow_coercion: bool = True) -> None:
761
- if _allow_coercion:
762
- # Field coercion logic (if any is required).
763
- if platform_name.startswith("urn:li:dataPlatform:"):
764
- platform_name = DataPlatformUrn.from_string(platform_name).platform_name
765
- platform_name = UrnEncoder.encode_string(platform_name)
766
-
767
- # Validation logic.
768
- if not platform_name:
769
- raise InvalidUrnError("DataPlatformUrn platform_name cannot be empty")
770
- if UrnEncoder.contains_reserved_char(platform_name):
771
- raise InvalidUrnError(f'DataPlatformUrn platform_name contains reserved characters')
772
-
773
- super().__init__(self.ENTITY_TYPE, [platform_name])
774
-
775
- @classmethod
776
- def _parse_ids(cls, entity_ids: List[str]) -> "DataPlatformUrn":
777
- if len(entity_ids) != cls._URN_PARTS:
778
- raise InvalidUrnError(f"DataPlatformUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
779
- return cls(platform_name=entity_ids[0], _allow_coercion=False)
780
-
781
- @classmethod
782
- def underlying_key_aspect_type(cls) -> Type["DataPlatformKeyClass"]:
783
- from datahub.metadata.schema_classes import DataPlatformKeyClass
784
-
785
- return DataPlatformKeyClass
786
-
787
- def to_key_aspect(self) -> "DataPlatformKeyClass":
788
- from datahub.metadata.schema_classes import DataPlatformKeyClass
789
-
790
- return DataPlatformKeyClass(platformName=self.platform_name)
791
-
792
- @classmethod
793
- def from_key_aspect(cls, key_aspect: "DataPlatformKeyClass") -> "DataPlatformUrn":
794
- return cls(platform_name=key_aspect.platformName)
795
-
796
- @classmethod
797
- @deprecated(reason="Use the constructor instead")
798
- def create_from_id(cls, id: str) -> "DataPlatformUrn":
799
- return cls(id)
800
-
801
- @property
802
- def platform_name(self) -> str:
803
- return self.entity_ids[0]
804
-
805
833
  if TYPE_CHECKING:
806
- from datahub.metadata.schema_classes import TagKeyClass
834
+ from datahub.metadata.schema_classes import DataHubConnectionKeyClass
807
835
 
808
- class TagUrn(_SpecificUrn):
809
- ENTITY_TYPE: ClassVar[str] = "tag"
836
+ class DataHubConnectionUrn(_SpecificUrn):
837
+ ENTITY_TYPE: ClassVar[str] = "dataHubConnection"
810
838
  _URN_PARTS: ClassVar[int] = 1
811
839
 
812
- def __init__(self, name: str, *, _allow_coercion: bool = True) -> None:
813
- if _allow_coercion:
814
- # Field coercion logic (if any is required).
815
- name = UrnEncoder.encode_string(name)
816
-
817
- # Validation logic.
818
- if not name:
819
- raise InvalidUrnError("TagUrn name cannot be empty")
820
- if UrnEncoder.contains_reserved_char(name):
821
- raise InvalidUrnError(f'TagUrn name contains reserved characters')
822
-
823
- super().__init__(self.ENTITY_TYPE, [name])
824
-
825
- @classmethod
826
- def _parse_ids(cls, entity_ids: List[str]) -> "TagUrn":
827
- if len(entity_ids) != cls._URN_PARTS:
828
- raise InvalidUrnError(f"TagUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
829
- return cls(name=entity_ids[0], _allow_coercion=False)
830
-
831
- @classmethod
832
- def underlying_key_aspect_type(cls) -> Type["TagKeyClass"]:
833
- from datahub.metadata.schema_classes import TagKeyClass
834
-
835
- return TagKeyClass
836
-
837
- def to_key_aspect(self) -> "TagKeyClass":
838
- from datahub.metadata.schema_classes import TagKeyClass
839
-
840
- return TagKeyClass(name=self.name)
841
-
842
- @classmethod
843
- def from_key_aspect(cls, key_aspect: "TagKeyClass") -> "TagUrn":
844
- return cls(name=key_aspect.name)
845
-
846
- @classmethod
847
- @deprecated(reason="Use the constructor instead")
848
- def create_from_id(cls, id: str) -> "TagUrn":
849
- return cls(id)
850
-
851
- @property
852
- def name(self) -> str:
853
- return self.entity_ids[0]
854
-
855
- if TYPE_CHECKING:
856
- from datahub.metadata.schema_classes import NotebookKeyClass
857
-
858
- class NotebookUrn(_SpecificUrn):
859
- ENTITY_TYPE: ClassVar[str] = "notebook"
860
- _URN_PARTS: ClassVar[int] = 2
861
-
862
- def __init__(self, notebook_tool: str, notebook_id: str, *, _allow_coercion: bool = True) -> None:
840
+ def __init__(self, id: Union["DataHubConnectionUrn", str], *, _allow_coercion: bool = True) -> None:
863
841
  if _allow_coercion:
864
842
  # Field coercion logic (if any is required).
865
- notebook_tool = UrnEncoder.encode_string(notebook_tool)
866
- notebook_id = UrnEncoder.encode_string(notebook_id)
843
+ if isinstance(id, str):
844
+ if id.startswith('urn:li:'):
845
+ try:
846
+ id = DataHubConnectionUrn.from_string(id)
847
+ except InvalidUrnError:
848
+ raise InvalidUrnError(f'Expecting a DataHubConnectionUrn but got {id}')
849
+ else:
850
+ id = UrnEncoder.encode_string(id)
867
851
 
868
852
  # Validation logic.
869
- if not notebook_tool:
870
- raise InvalidUrnError("NotebookUrn notebook_tool cannot be empty")
871
- if UrnEncoder.contains_reserved_char(notebook_tool):
872
- raise InvalidUrnError(f'NotebookUrn notebook_tool contains reserved characters')
873
- if not notebook_id:
874
- raise InvalidUrnError("NotebookUrn notebook_id cannot be empty")
875
- if UrnEncoder.contains_reserved_char(notebook_id):
876
- raise InvalidUrnError(f'NotebookUrn notebook_id contains reserved characters')
853
+ if not id:
854
+ raise InvalidUrnError("DataHubConnectionUrn id cannot be empty")
855
+ if isinstance(id, DataHubConnectionUrn):
856
+ id = id.id
857
+ elif isinstance(id, Urn):
858
+ raise InvalidUrnError(f'Expecting a DataHubConnectionUrn but got {id}')
859
+ if UrnEncoder.contains_reserved_char(id):
860
+ raise InvalidUrnError(f'DataHubConnectionUrn id contains reserved characters')
877
861
 
878
- super().__init__(self.ENTITY_TYPE, [notebook_tool, notebook_id])
862
+ super().__init__(self.ENTITY_TYPE, [id])
879
863
 
880
864
  @classmethod
881
- def _parse_ids(cls, entity_ids: List[str]) -> "NotebookUrn":
865
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubConnectionUrn":
882
866
  if len(entity_ids) != cls._URN_PARTS:
883
- raise InvalidUrnError(f"NotebookUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
884
- return cls(notebook_tool=entity_ids[0], notebook_id=entity_ids[1], _allow_coercion=False)
867
+ raise InvalidUrnError(f"DataHubConnectionUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
868
+ return cls(id=entity_ids[0], _allow_coercion=False)
885
869
 
886
870
  @classmethod
887
- def underlying_key_aspect_type(cls) -> Type["NotebookKeyClass"]:
888
- from datahub.metadata.schema_classes import NotebookKeyClass
871
+ def underlying_key_aspect_type(cls) -> Type["DataHubConnectionKeyClass"]:
872
+ from datahub.metadata.schema_classes import DataHubConnectionKeyClass
889
873
 
890
- return NotebookKeyClass
874
+ return DataHubConnectionKeyClass
891
875
 
892
- def to_key_aspect(self) -> "NotebookKeyClass":
893
- from datahub.metadata.schema_classes import NotebookKeyClass
876
+ def to_key_aspect(self) -> "DataHubConnectionKeyClass":
877
+ from datahub.metadata.schema_classes import DataHubConnectionKeyClass
894
878
 
895
- return NotebookKeyClass(notebookTool=self.notebook_tool, notebookId=self.notebook_id)
879
+ return DataHubConnectionKeyClass(id=self.id)
896
880
 
897
881
  @classmethod
898
- def from_key_aspect(cls, key_aspect: "NotebookKeyClass") -> "NotebookUrn":
899
- return cls(notebook_tool=key_aspect.notebookTool, notebook_id=key_aspect.notebookId)
900
-
901
- @deprecated(reason="Use .notebook_tool instead")
902
- def get_platform_id(self) -> str:
903
- return self.notebook_tool
904
-
905
- @deprecated(reason="Use .notebook_id instead")
906
- def get_notebook_id(self) -> str:
907
- return self.notebook_id
882
+ def from_key_aspect(cls, key_aspect: "DataHubConnectionKeyClass") -> "DataHubConnectionUrn":
883
+ return cls(id=key_aspect.id)
908
884
 
909
885
  @property
910
- def notebook_tool(self) -> str:
886
+ def id(self) -> str:
911
887
  return self.entity_ids[0]
912
888
 
913
- @property
914
- def notebook_id(self) -> str:
915
- return self.entity_ids[1]
916
-
917
889
  if TYPE_CHECKING:
918
- from datahub.metadata.schema_classes import DataHubAccessTokenKeyClass
890
+ from datahub.metadata.schema_classes import IncidentKeyClass
919
891
 
920
- class DataHubAccessTokenUrn(_SpecificUrn):
921
- ENTITY_TYPE: ClassVar[str] = "dataHubAccessToken"
892
+ class IncidentUrn(_SpecificUrn):
893
+ ENTITY_TYPE: ClassVar[str] = "incident"
922
894
  _URN_PARTS: ClassVar[int] = 1
923
895
 
924
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
896
+ def __init__(self, id: Union["IncidentUrn", str], *, _allow_coercion: bool = True) -> None:
925
897
  if _allow_coercion:
926
898
  # Field coercion logic (if any is required).
927
- id = UrnEncoder.encode_string(id)
899
+ if isinstance(id, str):
900
+ if id.startswith('urn:li:'):
901
+ try:
902
+ id = IncidentUrn.from_string(id)
903
+ except InvalidUrnError:
904
+ raise InvalidUrnError(f'Expecting a IncidentUrn but got {id}')
905
+ else:
906
+ id = UrnEncoder.encode_string(id)
928
907
 
929
908
  # Validation logic.
930
909
  if not id:
931
- raise InvalidUrnError("DataHubAccessTokenUrn id cannot be empty")
910
+ raise InvalidUrnError("IncidentUrn id cannot be empty")
911
+ if isinstance(id, IncidentUrn):
912
+ id = id.id
913
+ elif isinstance(id, Urn):
914
+ raise InvalidUrnError(f'Expecting a IncidentUrn but got {id}')
932
915
  if UrnEncoder.contains_reserved_char(id):
933
- raise InvalidUrnError(f'DataHubAccessTokenUrn id contains reserved characters')
916
+ raise InvalidUrnError(f'IncidentUrn id contains reserved characters')
934
917
 
935
918
  super().__init__(self.ENTITY_TYPE, [id])
936
919
 
937
920
  @classmethod
938
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubAccessTokenUrn":
921
+ def _parse_ids(cls, entity_ids: List[str]) -> "IncidentUrn":
939
922
  if len(entity_ids) != cls._URN_PARTS:
940
- raise InvalidUrnError(f"DataHubAccessTokenUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
923
+ raise InvalidUrnError(f"IncidentUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
941
924
  return cls(id=entity_ids[0], _allow_coercion=False)
942
925
 
943
926
  @classmethod
944
- def underlying_key_aspect_type(cls) -> Type["DataHubAccessTokenKeyClass"]:
945
- from datahub.metadata.schema_classes import DataHubAccessTokenKeyClass
927
+ def underlying_key_aspect_type(cls) -> Type["IncidentKeyClass"]:
928
+ from datahub.metadata.schema_classes import IncidentKeyClass
946
929
 
947
- return DataHubAccessTokenKeyClass
930
+ return IncidentKeyClass
948
931
 
949
- def to_key_aspect(self) -> "DataHubAccessTokenKeyClass":
950
- from datahub.metadata.schema_classes import DataHubAccessTokenKeyClass
932
+ def to_key_aspect(self) -> "IncidentKeyClass":
933
+ from datahub.metadata.schema_classes import IncidentKeyClass
951
934
 
952
- return DataHubAccessTokenKeyClass(id=self.id)
935
+ return IncidentKeyClass(id=self.id)
953
936
 
954
937
  @classmethod
955
- def from_key_aspect(cls, key_aspect: "DataHubAccessTokenKeyClass") -> "DataHubAccessTokenUrn":
938
+ def from_key_aspect(cls, key_aspect: "IncidentKeyClass") -> "IncidentUrn":
956
939
  return cls(id=key_aspect.id)
957
940
 
958
941
  @property
@@ -960,44 +943,55 @@ class DataHubAccessTokenUrn(_SpecificUrn):
960
943
  return self.entity_ids[0]
961
944
 
962
945
  if TYPE_CHECKING:
963
- from datahub.metadata.schema_classes import ERModelRelationshipKeyClass
946
+ from datahub.metadata.schema_classes import RoleKeyClass
964
947
 
965
- class ErModelRelationshipUrn(_SpecificUrn):
966
- ENTITY_TYPE: ClassVar[str] = "erModelRelationship"
948
+ class RoleUrn(_SpecificUrn):
949
+ ENTITY_TYPE: ClassVar[str] = "role"
967
950
  _URN_PARTS: ClassVar[int] = 1
968
951
 
969
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
952
+ def __init__(self, id: Union["RoleUrn", str], *, _allow_coercion: bool = True) -> None:
970
953
  if _allow_coercion:
971
954
  # Field coercion logic (if any is required).
972
- id = UrnEncoder.encode_string(id)
955
+ if isinstance(id, str):
956
+ if id.startswith('urn:li:'):
957
+ try:
958
+ id = RoleUrn.from_string(id)
959
+ except InvalidUrnError:
960
+ raise InvalidUrnError(f'Expecting a RoleUrn but got {id}')
961
+ else:
962
+ id = UrnEncoder.encode_string(id)
973
963
 
974
964
  # Validation logic.
975
965
  if not id:
976
- raise InvalidUrnError("ErModelRelationshipUrn id cannot be empty")
966
+ raise InvalidUrnError("RoleUrn id cannot be empty")
967
+ if isinstance(id, RoleUrn):
968
+ id = id.id
969
+ elif isinstance(id, Urn):
970
+ raise InvalidUrnError(f'Expecting a RoleUrn but got {id}')
977
971
  if UrnEncoder.contains_reserved_char(id):
978
- raise InvalidUrnError(f'ErModelRelationshipUrn id contains reserved characters')
972
+ raise InvalidUrnError(f'RoleUrn id contains reserved characters')
979
973
 
980
974
  super().__init__(self.ENTITY_TYPE, [id])
981
975
 
982
976
  @classmethod
983
- def _parse_ids(cls, entity_ids: List[str]) -> "ErModelRelationshipUrn":
977
+ def _parse_ids(cls, entity_ids: List[str]) -> "RoleUrn":
984
978
  if len(entity_ids) != cls._URN_PARTS:
985
- raise InvalidUrnError(f"ErModelRelationshipUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
979
+ raise InvalidUrnError(f"RoleUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
986
980
  return cls(id=entity_ids[0], _allow_coercion=False)
987
981
 
988
982
  @classmethod
989
- def underlying_key_aspect_type(cls) -> Type["ERModelRelationshipKeyClass"]:
990
- from datahub.metadata.schema_classes import ERModelRelationshipKeyClass
983
+ def underlying_key_aspect_type(cls) -> Type["RoleKeyClass"]:
984
+ from datahub.metadata.schema_classes import RoleKeyClass
991
985
 
992
- return ERModelRelationshipKeyClass
986
+ return RoleKeyClass
993
987
 
994
- def to_key_aspect(self) -> "ERModelRelationshipKeyClass":
995
- from datahub.metadata.schema_classes import ERModelRelationshipKeyClass
988
+ def to_key_aspect(self) -> "RoleKeyClass":
989
+ from datahub.metadata.schema_classes import RoleKeyClass
996
990
 
997
- return ERModelRelationshipKeyClass(id=self.id)
991
+ return RoleKeyClass(id=self.id)
998
992
 
999
993
  @classmethod
1000
- def from_key_aspect(cls, key_aspect: "ERModelRelationshipKeyClass") -> "ErModelRelationshipUrn":
994
+ def from_key_aspect(cls, key_aspect: "RoleKeyClass") -> "RoleUrn":
1001
995
  return cls(id=key_aspect.id)
1002
996
 
1003
997
  @property
@@ -1005,89 +999,116 @@ class ErModelRelationshipUrn(_SpecificUrn):
1005
999
  return self.entity_ids[0]
1006
1000
 
1007
1001
  if TYPE_CHECKING:
1008
- from datahub.metadata.schema_classes import OwnershipTypeKeyClass
1002
+ from datahub.metadata.schema_classes import DomainKeyClass
1009
1003
 
1010
- class OwnershipTypeUrn(_SpecificUrn):
1011
- ENTITY_TYPE: ClassVar[str] = "ownershipType"
1004
+ class DomainUrn(_SpecificUrn):
1005
+ ENTITY_TYPE: ClassVar[str] = "domain"
1012
1006
  _URN_PARTS: ClassVar[int] = 1
1013
1007
 
1014
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1008
+ def __init__(self, id: Union["DomainUrn", str], *, _allow_coercion: bool = True) -> None:
1015
1009
  if _allow_coercion:
1016
1010
  # Field coercion logic (if any is required).
1017
- id = UrnEncoder.encode_string(id)
1011
+ if isinstance(id, str):
1012
+ if id.startswith('urn:li:'):
1013
+ try:
1014
+ id = DomainUrn.from_string(id)
1015
+ except InvalidUrnError:
1016
+ raise InvalidUrnError(f'Expecting a DomainUrn but got {id}')
1017
+ else:
1018
+ id = UrnEncoder.encode_string(id)
1018
1019
 
1019
1020
  # Validation logic.
1020
1021
  if not id:
1021
- raise InvalidUrnError("OwnershipTypeUrn id cannot be empty")
1022
+ raise InvalidUrnError("DomainUrn id cannot be empty")
1023
+ if isinstance(id, DomainUrn):
1024
+ id = id.id
1025
+ elif isinstance(id, Urn):
1026
+ raise InvalidUrnError(f'Expecting a DomainUrn but got {id}')
1022
1027
  if UrnEncoder.contains_reserved_char(id):
1023
- raise InvalidUrnError(f'OwnershipTypeUrn id contains reserved characters')
1028
+ raise InvalidUrnError(f'DomainUrn id contains reserved characters')
1024
1029
 
1025
1030
  super().__init__(self.ENTITY_TYPE, [id])
1026
1031
 
1027
1032
  @classmethod
1028
- def _parse_ids(cls, entity_ids: List[str]) -> "OwnershipTypeUrn":
1033
+ def _parse_ids(cls, entity_ids: List[str]) -> "DomainUrn":
1029
1034
  if len(entity_ids) != cls._URN_PARTS:
1030
- raise InvalidUrnError(f"OwnershipTypeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1035
+ raise InvalidUrnError(f"DomainUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1031
1036
  return cls(id=entity_ids[0], _allow_coercion=False)
1032
1037
 
1033
1038
  @classmethod
1034
- def underlying_key_aspect_type(cls) -> Type["OwnershipTypeKeyClass"]:
1035
- from datahub.metadata.schema_classes import OwnershipTypeKeyClass
1039
+ def underlying_key_aspect_type(cls) -> Type["DomainKeyClass"]:
1040
+ from datahub.metadata.schema_classes import DomainKeyClass
1036
1041
 
1037
- return OwnershipTypeKeyClass
1042
+ return DomainKeyClass
1038
1043
 
1039
- def to_key_aspect(self) -> "OwnershipTypeKeyClass":
1040
- from datahub.metadata.schema_classes import OwnershipTypeKeyClass
1044
+ def to_key_aspect(self) -> "DomainKeyClass":
1045
+ from datahub.metadata.schema_classes import DomainKeyClass
1041
1046
 
1042
- return OwnershipTypeKeyClass(id=self.id)
1047
+ return DomainKeyClass(id=self.id)
1043
1048
 
1044
1049
  @classmethod
1045
- def from_key_aspect(cls, key_aspect: "OwnershipTypeKeyClass") -> "OwnershipTypeUrn":
1050
+ def from_key_aspect(cls, key_aspect: "DomainKeyClass") -> "DomainUrn":
1046
1051
  return cls(id=key_aspect.id)
1047
1052
 
1053
+ @classmethod
1054
+ @deprecated(reason="Use the constructor instead")
1055
+ def create_from_id(cls, id: str) -> "DomainUrn":
1056
+ return cls(id)
1057
+
1048
1058
  @property
1049
1059
  def id(self) -> str:
1050
1060
  return self.entity_ids[0]
1051
1061
 
1052
1062
  if TYPE_CHECKING:
1053
- from datahub.metadata.schema_classes import QueryKeyClass
1063
+ from datahub.metadata.schema_classes import DataContractKeyClass
1054
1064
 
1055
- class QueryUrn(_SpecificUrn):
1056
- ENTITY_TYPE: ClassVar[str] = "query"
1065
+ class DataContractUrn(_SpecificUrn):
1066
+ ENTITY_TYPE: ClassVar[str] = "dataContract"
1057
1067
  _URN_PARTS: ClassVar[int] = 1
1058
1068
 
1059
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1069
+ def __init__(self, id: Union["DataContractUrn", str], *, _allow_coercion: bool = True) -> None:
1060
1070
  if _allow_coercion:
1061
1071
  # Field coercion logic (if any is required).
1062
- id = UrnEncoder.encode_string(id)
1072
+ if isinstance(id, str):
1073
+ if id.startswith('urn:li:'):
1074
+ try:
1075
+ id = DataContractUrn.from_string(id)
1076
+ except InvalidUrnError:
1077
+ raise InvalidUrnError(f'Expecting a DataContractUrn but got {id}')
1078
+ else:
1079
+ id = UrnEncoder.encode_string(id)
1063
1080
 
1064
1081
  # Validation logic.
1065
1082
  if not id:
1066
- raise InvalidUrnError("QueryUrn id cannot be empty")
1083
+ raise InvalidUrnError("DataContractUrn id cannot be empty")
1084
+ if isinstance(id, DataContractUrn):
1085
+ id = id.id
1086
+ elif isinstance(id, Urn):
1087
+ raise InvalidUrnError(f'Expecting a DataContractUrn but got {id}')
1067
1088
  if UrnEncoder.contains_reserved_char(id):
1068
- raise InvalidUrnError(f'QueryUrn id contains reserved characters')
1089
+ raise InvalidUrnError(f'DataContractUrn id contains reserved characters')
1069
1090
 
1070
1091
  super().__init__(self.ENTITY_TYPE, [id])
1071
1092
 
1072
1093
  @classmethod
1073
- def _parse_ids(cls, entity_ids: List[str]) -> "QueryUrn":
1094
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataContractUrn":
1074
1095
  if len(entity_ids) != cls._URN_PARTS:
1075
- raise InvalidUrnError(f"QueryUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1096
+ raise InvalidUrnError(f"DataContractUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1076
1097
  return cls(id=entity_ids[0], _allow_coercion=False)
1077
1098
 
1078
1099
  @classmethod
1079
- def underlying_key_aspect_type(cls) -> Type["QueryKeyClass"]:
1080
- from datahub.metadata.schema_classes import QueryKeyClass
1100
+ def underlying_key_aspect_type(cls) -> Type["DataContractKeyClass"]:
1101
+ from datahub.metadata.schema_classes import DataContractKeyClass
1081
1102
 
1082
- return QueryKeyClass
1103
+ return DataContractKeyClass
1083
1104
 
1084
- def to_key_aspect(self) -> "QueryKeyClass":
1085
- from datahub.metadata.schema_classes import QueryKeyClass
1105
+ def to_key_aspect(self) -> "DataContractKeyClass":
1106
+ from datahub.metadata.schema_classes import DataContractKeyClass
1086
1107
 
1087
- return QueryKeyClass(id=self.id)
1108
+ return DataContractKeyClass(id=self.id)
1088
1109
 
1089
1110
  @classmethod
1090
- def from_key_aspect(cls, key_aspect: "QueryKeyClass") -> "QueryUrn":
1111
+ def from_key_aspect(cls, key_aspect: "DataContractKeyClass") -> "DataContractUrn":
1091
1112
  return cls(id=key_aspect.id)
1092
1113
 
1093
1114
  @property
@@ -1101,14 +1122,25 @@ class AssertionUrn(_SpecificUrn):
1101
1122
  ENTITY_TYPE: ClassVar[str] = "assertion"
1102
1123
  _URN_PARTS: ClassVar[int] = 1
1103
1124
 
1104
- def __init__(self, assertion_id: str, *, _allow_coercion: bool = True) -> None:
1125
+ def __init__(self, assertion_id: Union["AssertionUrn", str], *, _allow_coercion: bool = True) -> None:
1105
1126
  if _allow_coercion:
1106
1127
  # Field coercion logic (if any is required).
1107
- assertion_id = UrnEncoder.encode_string(assertion_id)
1128
+ if isinstance(assertion_id, str):
1129
+ if assertion_id.startswith('urn:li:'):
1130
+ try:
1131
+ assertion_id = AssertionUrn.from_string(assertion_id)
1132
+ except InvalidUrnError:
1133
+ raise InvalidUrnError(f'Expecting a AssertionUrn but got {assertion_id}')
1134
+ else:
1135
+ assertion_id = UrnEncoder.encode_string(assertion_id)
1108
1136
 
1109
1137
  # Validation logic.
1110
1138
  if not assertion_id:
1111
1139
  raise InvalidUrnError("AssertionUrn assertion_id cannot be empty")
1140
+ if isinstance(assertion_id, AssertionUrn):
1141
+ assertion_id = assertion_id.assertion_id
1142
+ elif isinstance(assertion_id, Urn):
1143
+ raise InvalidUrnError(f'Expecting a AssertionUrn but got {assertion_id}')
1112
1144
  if UrnEncoder.contains_reserved_char(assertion_id):
1113
1145
  raise InvalidUrnError(f'AssertionUrn assertion_id contains reserved characters')
1114
1146
 
@@ -1140,44 +1172,219 @@ class AssertionUrn(_SpecificUrn):
1140
1172
  return self.entity_ids[0]
1141
1173
 
1142
1174
  if TYPE_CHECKING:
1143
- from datahub.metadata.schema_classes import DataContractKeyClass
1175
+ from datahub.metadata.schema_classes import DatasetKeyClass
1144
1176
 
1145
- class DataContractUrn(_SpecificUrn):
1146
- ENTITY_TYPE: ClassVar[str] = "dataContract"
1177
+ class DatasetUrn(_SpecificUrn):
1178
+ ENTITY_TYPE: ClassVar[str] = "dataset"
1179
+ _URN_PARTS: ClassVar[int] = 3
1180
+
1181
+ def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
1182
+ if _allow_coercion:
1183
+ # Field coercion logic (if any is required).
1184
+ platform = DataPlatformUrn(platform).urn()
1185
+ name = UrnEncoder.encode_string(name)
1186
+ env = env.upper()
1187
+
1188
+ # Validation logic.
1189
+ if not platform:
1190
+ raise InvalidUrnError("DatasetUrn platform cannot be empty")
1191
+ platform = str(platform) # convert urn type to str
1192
+ assert DataPlatformUrn.from_string(platform)
1193
+ if not name:
1194
+ raise InvalidUrnError("DatasetUrn name cannot be empty")
1195
+ if UrnEncoder.contains_reserved_char(name):
1196
+ raise InvalidUrnError(f'DatasetUrn name contains reserved characters')
1197
+ if not env:
1198
+ raise InvalidUrnError("DatasetUrn env cannot be empty")
1199
+ if UrnEncoder.contains_reserved_char(env):
1200
+ raise InvalidUrnError(f'DatasetUrn env contains reserved characters')
1201
+
1202
+ super().__init__(self.ENTITY_TYPE, [platform, name, env])
1203
+
1204
+ @classmethod
1205
+ def _parse_ids(cls, entity_ids: List[str]) -> "DatasetUrn":
1206
+ if len(entity_ids) != cls._URN_PARTS:
1207
+ raise InvalidUrnError(f"DatasetUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1208
+ return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
1209
+
1210
+ @classmethod
1211
+ def underlying_key_aspect_type(cls) -> Type["DatasetKeyClass"]:
1212
+ from datahub.metadata.schema_classes import DatasetKeyClass
1213
+
1214
+ return DatasetKeyClass
1215
+
1216
+ def to_key_aspect(self) -> "DatasetKeyClass":
1217
+ from datahub.metadata.schema_classes import DatasetKeyClass
1218
+
1219
+ return DatasetKeyClass(platform=self.platform, name=self.name, origin=self.env)
1220
+
1221
+ @classmethod
1222
+ def from_key_aspect(cls, key_aspect: "DatasetKeyClass") -> "DatasetUrn":
1223
+ return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
1224
+
1225
+ @classmethod
1226
+ def create_from_ids(
1227
+ cls,
1228
+ platform_id: str,
1229
+ table_name: str,
1230
+ env: str,
1231
+ platform_instance: Optional[str] = None,
1232
+ ) -> "DatasetUrn":
1233
+ return DatasetUrn(
1234
+ platform=platform_id,
1235
+ name=f"{platform_instance}.{table_name}" if platform_instance else table_name,
1236
+ env=env,
1237
+ )
1238
+
1239
+ from datahub.utilities.urns.field_paths import get_simple_field_path_from_v2_field_path as _get_simple_field_path_from_v2_field_path
1240
+
1241
+ get_simple_field_path_from_v2_field_path = staticmethod(deprecated(reason='Use the function from the field_paths module instead')(_get_simple_field_path_from_v2_field_path))
1242
+
1243
+ def get_data_platform_urn(self) -> "DataPlatformUrn":
1244
+ return DataPlatformUrn.from_string(self.platform)
1245
+
1246
+ @deprecated(reason="Use .name instead")
1247
+ def get_dataset_name(self) -> str:
1248
+ return self.name
1249
+
1250
+ @deprecated(reason="Use .env instead")
1251
+ def get_env(self) -> str:
1252
+ return self.env
1253
+
1254
+ @property
1255
+ def platform(self) -> str:
1256
+ return self.entity_ids[0]
1257
+
1258
+ @property
1259
+ def name(self) -> str:
1260
+ return self.entity_ids[1]
1261
+
1262
+ @property
1263
+ def env(self) -> str:
1264
+ return self.entity_ids[2]
1265
+
1266
+ if TYPE_CHECKING:
1267
+ from datahub.metadata.schema_classes import DataJobKeyClass
1268
+
1269
+ class DataJobUrn(_SpecificUrn):
1270
+ ENTITY_TYPE: ClassVar[str] = "dataJob"
1271
+ _URN_PARTS: ClassVar[int] = 2
1272
+
1273
+ def __init__(self, flow: Union["DataFlowUrn", str], job_id: str, *, _allow_coercion: bool = True) -> None:
1274
+ if _allow_coercion:
1275
+ # Field coercion logic (if any is required).
1276
+ if isinstance(flow, str):
1277
+ if flow.startswith('urn:li:'):
1278
+ try:
1279
+ flow = DataFlowUrn.from_string(flow)
1280
+ except InvalidUrnError:
1281
+ raise InvalidUrnError(f'Expecting a DataFlowUrn but got {flow}')
1282
+ else:
1283
+ flow = UrnEncoder.encode_string(flow)
1284
+ job_id = UrnEncoder.encode_string(job_id)
1285
+
1286
+ # Validation logic.
1287
+ if not flow:
1288
+ raise InvalidUrnError("DataJobUrn flow cannot be empty")
1289
+ flow = str(flow) # convert urn type to str
1290
+ assert DataFlowUrn.from_string(flow)
1291
+ if not job_id:
1292
+ raise InvalidUrnError("DataJobUrn job_id cannot be empty")
1293
+ if UrnEncoder.contains_reserved_char(job_id):
1294
+ raise InvalidUrnError(f'DataJobUrn job_id contains reserved characters')
1295
+
1296
+ super().__init__(self.ENTITY_TYPE, [flow, job_id])
1297
+
1298
+ @classmethod
1299
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataJobUrn":
1300
+ if len(entity_ids) != cls._URN_PARTS:
1301
+ raise InvalidUrnError(f"DataJobUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1302
+ return cls(flow=entity_ids[0], job_id=entity_ids[1], _allow_coercion=False)
1303
+
1304
+ @classmethod
1305
+ def underlying_key_aspect_type(cls) -> Type["DataJobKeyClass"]:
1306
+ from datahub.metadata.schema_classes import DataJobKeyClass
1307
+
1308
+ return DataJobKeyClass
1309
+
1310
+ def to_key_aspect(self) -> "DataJobKeyClass":
1311
+ from datahub.metadata.schema_classes import DataJobKeyClass
1312
+
1313
+ return DataJobKeyClass(flow=self.flow, jobId=self.job_id)
1314
+
1315
+ @classmethod
1316
+ def from_key_aspect(cls, key_aspect: "DataJobKeyClass") -> "DataJobUrn":
1317
+ return cls(flow=key_aspect.flow, job_id=key_aspect.jobId)
1318
+
1319
+ @classmethod
1320
+ def create_from_ids(cls, data_flow_urn: str, job_id: str) -> "DataJobUrn":
1321
+ return cls(data_flow_urn, job_id)
1322
+
1323
+ def get_data_flow_urn(self) -> "DataFlowUrn":
1324
+ return DataFlowUrn.from_string(self.flow)
1325
+
1326
+ @deprecated(reason="Use .job_id instead")
1327
+ def get_job_id(self) -> str:
1328
+ return self.job_id
1329
+
1330
+ @property
1331
+ def flow(self) -> str:
1332
+ return self.entity_ids[0]
1333
+
1334
+ @property
1335
+ def job_id(self) -> str:
1336
+ return self.entity_ids[1]
1337
+
1338
+ if TYPE_CHECKING:
1339
+ from datahub.metadata.schema_classes import DataHubSecretKeyClass
1340
+
1341
+ class DataHubSecretUrn(_SpecificUrn):
1342
+ ENTITY_TYPE: ClassVar[str] = "dataHubSecret"
1147
1343
  _URN_PARTS: ClassVar[int] = 1
1148
1344
 
1149
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1345
+ def __init__(self, id: Union["DataHubSecretUrn", str], *, _allow_coercion: bool = True) -> None:
1150
1346
  if _allow_coercion:
1151
1347
  # Field coercion logic (if any is required).
1152
- id = UrnEncoder.encode_string(id)
1348
+ if isinstance(id, str):
1349
+ if id.startswith('urn:li:'):
1350
+ try:
1351
+ id = DataHubSecretUrn.from_string(id)
1352
+ except InvalidUrnError:
1353
+ raise InvalidUrnError(f'Expecting a DataHubSecretUrn but got {id}')
1354
+ else:
1355
+ id = UrnEncoder.encode_string(id)
1153
1356
 
1154
1357
  # Validation logic.
1155
1358
  if not id:
1156
- raise InvalidUrnError("DataContractUrn id cannot be empty")
1359
+ raise InvalidUrnError("DataHubSecretUrn id cannot be empty")
1360
+ if isinstance(id, DataHubSecretUrn):
1361
+ id = id.id
1362
+ elif isinstance(id, Urn):
1363
+ raise InvalidUrnError(f'Expecting a DataHubSecretUrn but got {id}')
1157
1364
  if UrnEncoder.contains_reserved_char(id):
1158
- raise InvalidUrnError(f'DataContractUrn id contains reserved characters')
1365
+ raise InvalidUrnError(f'DataHubSecretUrn id contains reserved characters')
1159
1366
 
1160
1367
  super().__init__(self.ENTITY_TYPE, [id])
1161
1368
 
1162
1369
  @classmethod
1163
- def _parse_ids(cls, entity_ids: List[str]) -> "DataContractUrn":
1370
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubSecretUrn":
1164
1371
  if len(entity_ids) != cls._URN_PARTS:
1165
- raise InvalidUrnError(f"DataContractUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1372
+ raise InvalidUrnError(f"DataHubSecretUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1166
1373
  return cls(id=entity_ids[0], _allow_coercion=False)
1167
1374
 
1168
1375
  @classmethod
1169
- def underlying_key_aspect_type(cls) -> Type["DataContractKeyClass"]:
1170
- from datahub.metadata.schema_classes import DataContractKeyClass
1376
+ def underlying_key_aspect_type(cls) -> Type["DataHubSecretKeyClass"]:
1377
+ from datahub.metadata.schema_classes import DataHubSecretKeyClass
1171
1378
 
1172
- return DataContractKeyClass
1379
+ return DataHubSecretKeyClass
1173
1380
 
1174
- def to_key_aspect(self) -> "DataContractKeyClass":
1175
- from datahub.metadata.schema_classes import DataContractKeyClass
1381
+ def to_key_aspect(self) -> "DataHubSecretKeyClass":
1382
+ from datahub.metadata.schema_classes import DataHubSecretKeyClass
1176
1383
 
1177
- return DataContractKeyClass(id=self.id)
1384
+ return DataHubSecretKeyClass(id=self.id)
1178
1385
 
1179
1386
  @classmethod
1180
- def from_key_aspect(cls, key_aspect: "DataContractKeyClass") -> "DataContractUrn":
1387
+ def from_key_aspect(cls, key_aspect: "DataHubSecretKeyClass") -> "DataHubSecretUrn":
1181
1388
  return cls(id=key_aspect.id)
1182
1389
 
1183
1390
  @property
@@ -1185,44 +1392,55 @@ class DataContractUrn(_SpecificUrn):
1185
1392
  return self.entity_ids[0]
1186
1393
 
1187
1394
  if TYPE_CHECKING:
1188
- from datahub.metadata.schema_classes import DataHubPolicyKeyClass
1395
+ from datahub.metadata.schema_classes import OwnershipTypeKeyClass
1189
1396
 
1190
- class DataHubPolicyUrn(_SpecificUrn):
1191
- ENTITY_TYPE: ClassVar[str] = "dataHubPolicy"
1397
+ class OwnershipTypeUrn(_SpecificUrn):
1398
+ ENTITY_TYPE: ClassVar[str] = "ownershipType"
1192
1399
  _URN_PARTS: ClassVar[int] = 1
1193
1400
 
1194
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1401
+ def __init__(self, id: Union["OwnershipTypeUrn", str], *, _allow_coercion: bool = True) -> None:
1195
1402
  if _allow_coercion:
1196
1403
  # Field coercion logic (if any is required).
1197
- id = UrnEncoder.encode_string(id)
1404
+ if isinstance(id, str):
1405
+ if id.startswith('urn:li:'):
1406
+ try:
1407
+ id = OwnershipTypeUrn.from_string(id)
1408
+ except InvalidUrnError:
1409
+ raise InvalidUrnError(f'Expecting a OwnershipTypeUrn but got {id}')
1410
+ else:
1411
+ id = UrnEncoder.encode_string(id)
1198
1412
 
1199
1413
  # Validation logic.
1200
1414
  if not id:
1201
- raise InvalidUrnError("DataHubPolicyUrn id cannot be empty")
1415
+ raise InvalidUrnError("OwnershipTypeUrn id cannot be empty")
1416
+ if isinstance(id, OwnershipTypeUrn):
1417
+ id = id.id
1418
+ elif isinstance(id, Urn):
1419
+ raise InvalidUrnError(f'Expecting a OwnershipTypeUrn but got {id}')
1202
1420
  if UrnEncoder.contains_reserved_char(id):
1203
- raise InvalidUrnError(f'DataHubPolicyUrn id contains reserved characters')
1421
+ raise InvalidUrnError(f'OwnershipTypeUrn id contains reserved characters')
1204
1422
 
1205
1423
  super().__init__(self.ENTITY_TYPE, [id])
1206
1424
 
1207
1425
  @classmethod
1208
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubPolicyUrn":
1426
+ def _parse_ids(cls, entity_ids: List[str]) -> "OwnershipTypeUrn":
1209
1427
  if len(entity_ids) != cls._URN_PARTS:
1210
- raise InvalidUrnError(f"DataHubPolicyUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1428
+ raise InvalidUrnError(f"OwnershipTypeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1211
1429
  return cls(id=entity_ids[0], _allow_coercion=False)
1212
1430
 
1213
1431
  @classmethod
1214
- def underlying_key_aspect_type(cls) -> Type["DataHubPolicyKeyClass"]:
1215
- from datahub.metadata.schema_classes import DataHubPolicyKeyClass
1432
+ def underlying_key_aspect_type(cls) -> Type["OwnershipTypeKeyClass"]:
1433
+ from datahub.metadata.schema_classes import OwnershipTypeKeyClass
1216
1434
 
1217
- return DataHubPolicyKeyClass
1435
+ return OwnershipTypeKeyClass
1218
1436
 
1219
- def to_key_aspect(self) -> "DataHubPolicyKeyClass":
1220
- from datahub.metadata.schema_classes import DataHubPolicyKeyClass
1437
+ def to_key_aspect(self) -> "OwnershipTypeKeyClass":
1438
+ from datahub.metadata.schema_classes import OwnershipTypeKeyClass
1221
1439
 
1222
- return DataHubPolicyKeyClass(id=self.id)
1440
+ return OwnershipTypeKeyClass(id=self.id)
1223
1441
 
1224
1442
  @classmethod
1225
- def from_key_aspect(cls, key_aspect: "DataHubPolicyKeyClass") -> "DataHubPolicyUrn":
1443
+ def from_key_aspect(cls, key_aspect: "OwnershipTypeKeyClass") -> "OwnershipTypeUrn":
1226
1444
  return cls(id=key_aspect.id)
1227
1445
 
1228
1446
  @property
@@ -1230,297 +1448,334 @@ class DataHubPolicyUrn(_SpecificUrn):
1230
1448
  return self.entity_ids[0]
1231
1449
 
1232
1450
  if TYPE_CHECKING:
1233
- from datahub.metadata.schema_classes import DataHubActionKeyClass
1451
+ from datahub.metadata.schema_classes import SchemaFieldKeyClass
1234
1452
 
1235
- class DataHubActionUrn(_SpecificUrn):
1236
- ENTITY_TYPE: ClassVar[str] = "dataHubAction"
1237
- _URN_PARTS: ClassVar[int] = 1
1453
+ class SchemaFieldUrn(_SpecificUrn):
1454
+ ENTITY_TYPE: ClassVar[str] = "schemaField"
1455
+ _URN_PARTS: ClassVar[int] = 2
1238
1456
 
1239
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1457
+ def __init__(self, parent: Union["Urn", str], field_path: str, *, _allow_coercion: bool = True) -> None:
1240
1458
  if _allow_coercion:
1241
1459
  # Field coercion logic (if any is required).
1242
- id = UrnEncoder.encode_string(id)
1460
+ if isinstance(parent, str):
1461
+ if parent.startswith('urn:li:'):
1462
+ try:
1463
+ parent = Urn.from_string(parent)
1464
+ except InvalidUrnError:
1465
+ raise InvalidUrnError(f'Expecting a Urn but got {parent}')
1466
+ else:
1467
+ parent = UrnEncoder.encode_string(parent)
1468
+ field_path = UrnEncoder.encode_string(field_path)
1243
1469
 
1244
1470
  # Validation logic.
1245
- if not id:
1246
- raise InvalidUrnError("DataHubActionUrn id cannot be empty")
1247
- if UrnEncoder.contains_reserved_char(id):
1248
- raise InvalidUrnError(f'DataHubActionUrn id contains reserved characters')
1471
+ if not parent:
1472
+ raise InvalidUrnError("SchemaFieldUrn parent cannot be empty")
1473
+ parent = str(parent) # convert urn type to str
1474
+ assert Urn.from_string(parent)
1475
+ if not field_path:
1476
+ raise InvalidUrnError("SchemaFieldUrn field_path cannot be empty")
1477
+ if UrnEncoder.contains_reserved_char(field_path):
1478
+ raise InvalidUrnError(f'SchemaFieldUrn field_path contains reserved characters')
1249
1479
 
1250
- super().__init__(self.ENTITY_TYPE, [id])
1480
+ super().__init__(self.ENTITY_TYPE, [parent, field_path])
1251
1481
 
1252
1482
  @classmethod
1253
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubActionUrn":
1483
+ def _parse_ids(cls, entity_ids: List[str]) -> "SchemaFieldUrn":
1254
1484
  if len(entity_ids) != cls._URN_PARTS:
1255
- raise InvalidUrnError(f"DataHubActionUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1256
- return cls(id=entity_ids[0], _allow_coercion=False)
1485
+ raise InvalidUrnError(f"SchemaFieldUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1486
+ return cls(parent=entity_ids[0], field_path=entity_ids[1], _allow_coercion=False)
1257
1487
 
1258
1488
  @classmethod
1259
- def underlying_key_aspect_type(cls) -> Type["DataHubActionKeyClass"]:
1260
- from datahub.metadata.schema_classes import DataHubActionKeyClass
1489
+ def underlying_key_aspect_type(cls) -> Type["SchemaFieldKeyClass"]:
1490
+ from datahub.metadata.schema_classes import SchemaFieldKeyClass
1261
1491
 
1262
- return DataHubActionKeyClass
1492
+ return SchemaFieldKeyClass
1263
1493
 
1264
- def to_key_aspect(self) -> "DataHubActionKeyClass":
1265
- from datahub.metadata.schema_classes import DataHubActionKeyClass
1494
+ def to_key_aspect(self) -> "SchemaFieldKeyClass":
1495
+ from datahub.metadata.schema_classes import SchemaFieldKeyClass
1266
1496
 
1267
- return DataHubActionKeyClass(id=self.id)
1497
+ return SchemaFieldKeyClass(parent=self.parent, fieldPath=self.field_path)
1268
1498
 
1269
1499
  @classmethod
1270
- def from_key_aspect(cls, key_aspect: "DataHubActionKeyClass") -> "DataHubActionUrn":
1271
- return cls(id=key_aspect.id)
1500
+ def from_key_aspect(cls, key_aspect: "SchemaFieldKeyClass") -> "SchemaFieldUrn":
1501
+ return cls(parent=key_aspect.parent, field_path=key_aspect.fieldPath)
1272
1502
 
1273
1503
  @property
1274
- def id(self) -> str:
1504
+ def parent(self) -> str:
1275
1505
  return self.entity_ids[0]
1276
1506
 
1507
+ @property
1508
+ def field_path(self) -> str:
1509
+ return self.entity_ids[1]
1510
+
1277
1511
  if TYPE_CHECKING:
1278
- from datahub.metadata.schema_classes import ContainerKeyClass
1512
+ from datahub.metadata.schema_classes import MLFeatureTableKeyClass
1279
1513
 
1280
- class ContainerUrn(_SpecificUrn):
1281
- ENTITY_TYPE: ClassVar[str] = "container"
1282
- _URN_PARTS: ClassVar[int] = 1
1514
+ class MlFeatureTableUrn(_SpecificUrn):
1515
+ ENTITY_TYPE: ClassVar[str] = "mlFeatureTable"
1516
+ _URN_PARTS: ClassVar[int] = 2
1283
1517
 
1284
- def __init__(self, guid: str, *, _allow_coercion: bool = True) -> None:
1518
+ def __init__(self, platform: Union["DataPlatformUrn", str], name: str, *, _allow_coercion: bool = True) -> None:
1285
1519
  if _allow_coercion:
1286
1520
  # Field coercion logic (if any is required).
1287
- guid = UrnEncoder.encode_string(guid)
1521
+ platform = DataPlatformUrn(platform).urn()
1522
+ name = UrnEncoder.encode_string(name)
1288
1523
 
1289
1524
  # Validation logic.
1290
- if not guid:
1291
- raise InvalidUrnError("ContainerUrn guid cannot be empty")
1292
- if UrnEncoder.contains_reserved_char(guid):
1293
- raise InvalidUrnError(f'ContainerUrn guid contains reserved characters')
1525
+ if not platform:
1526
+ raise InvalidUrnError("MlFeatureTableUrn platform cannot be empty")
1527
+ platform = str(platform) # convert urn type to str
1528
+ assert DataPlatformUrn.from_string(platform)
1529
+ if not name:
1530
+ raise InvalidUrnError("MlFeatureTableUrn name cannot be empty")
1531
+ if UrnEncoder.contains_reserved_char(name):
1532
+ raise InvalidUrnError(f'MlFeatureTableUrn name contains reserved characters')
1294
1533
 
1295
- super().__init__(self.ENTITY_TYPE, [guid])
1534
+ super().__init__(self.ENTITY_TYPE, [platform, name])
1296
1535
 
1297
1536
  @classmethod
1298
- def _parse_ids(cls, entity_ids: List[str]) -> "ContainerUrn":
1537
+ def _parse_ids(cls, entity_ids: List[str]) -> "MlFeatureTableUrn":
1299
1538
  if len(entity_ids) != cls._URN_PARTS:
1300
- raise InvalidUrnError(f"ContainerUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1301
- return cls(guid=entity_ids[0], _allow_coercion=False)
1539
+ raise InvalidUrnError(f"MlFeatureTableUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1540
+ return cls(platform=entity_ids[0], name=entity_ids[1], _allow_coercion=False)
1302
1541
 
1303
1542
  @classmethod
1304
- def underlying_key_aspect_type(cls) -> Type["ContainerKeyClass"]:
1305
- from datahub.metadata.schema_classes import ContainerKeyClass
1543
+ def underlying_key_aspect_type(cls) -> Type["MLFeatureTableKeyClass"]:
1544
+ from datahub.metadata.schema_classes import MLFeatureTableKeyClass
1306
1545
 
1307
- return ContainerKeyClass
1546
+ return MLFeatureTableKeyClass
1308
1547
 
1309
- def to_key_aspect(self) -> "ContainerKeyClass":
1310
- from datahub.metadata.schema_classes import ContainerKeyClass
1548
+ def to_key_aspect(self) -> "MLFeatureTableKeyClass":
1549
+ from datahub.metadata.schema_classes import MLFeatureTableKeyClass
1311
1550
 
1312
- return ContainerKeyClass(guid=self.guid)
1551
+ return MLFeatureTableKeyClass(platform=self.platform, name=self.name)
1313
1552
 
1314
1553
  @classmethod
1315
- def from_key_aspect(cls, key_aspect: "ContainerKeyClass") -> "ContainerUrn":
1316
- return cls(guid=key_aspect.guid)
1554
+ def from_key_aspect(cls, key_aspect: "MLFeatureTableKeyClass") -> "MlFeatureTableUrn":
1555
+ return cls(platform=key_aspect.platform, name=key_aspect.name)
1317
1556
 
1318
1557
  @property
1319
- def guid(self) -> str:
1558
+ def platform(self) -> str:
1320
1559
  return self.entity_ids[0]
1321
1560
 
1561
+ @property
1562
+ def name(self) -> str:
1563
+ return self.entity_ids[1]
1564
+
1322
1565
  if TYPE_CHECKING:
1323
- from datahub.metadata.schema_classes import DataProcessKeyClass
1566
+ from datahub.metadata.schema_classes import DashboardKeyClass
1324
1567
 
1325
- class DataProcessUrn(_SpecificUrn):
1326
- ENTITY_TYPE: ClassVar[str] = "dataProcess"
1327
- _URN_PARTS: ClassVar[int] = 3
1568
+ class DashboardUrn(_SpecificUrn):
1569
+ ENTITY_TYPE: ClassVar[str] = "dashboard"
1570
+ _URN_PARTS: ClassVar[int] = 2
1328
1571
 
1329
- def __init__(self, name: str, orchestrator: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
1572
+ def __init__(self, dashboard_tool: str, dashboard_id: str, *, _allow_coercion: bool = True) -> None:
1330
1573
  if _allow_coercion:
1331
1574
  # Field coercion logic (if any is required).
1332
- name = UrnEncoder.encode_string(name)
1333
- orchestrator = UrnEncoder.encode_string(orchestrator)
1334
- env = env.upper()
1335
- env = UrnEncoder.encode_string(env)
1575
+ dashboard_tool = UrnEncoder.encode_string(dashboard_tool)
1576
+ dashboard_id = UrnEncoder.encode_string(dashboard_id)
1336
1577
 
1337
1578
  # Validation logic.
1338
- if not name:
1339
- raise InvalidUrnError("DataProcessUrn name cannot be empty")
1340
- if UrnEncoder.contains_reserved_char(name):
1341
- raise InvalidUrnError(f'DataProcessUrn name contains reserved characters')
1342
- if not orchestrator:
1343
- raise InvalidUrnError("DataProcessUrn orchestrator cannot be empty")
1344
- if UrnEncoder.contains_reserved_char(orchestrator):
1345
- raise InvalidUrnError(f'DataProcessUrn orchestrator contains reserved characters')
1346
- if not env:
1347
- raise InvalidUrnError("DataProcessUrn env cannot be empty")
1348
- if UrnEncoder.contains_reserved_char(env):
1349
- raise InvalidUrnError(f'DataProcessUrn env contains reserved characters')
1579
+ if not dashboard_tool:
1580
+ raise InvalidUrnError("DashboardUrn dashboard_tool cannot be empty")
1581
+ if UrnEncoder.contains_reserved_char(dashboard_tool):
1582
+ raise InvalidUrnError(f'DashboardUrn dashboard_tool contains reserved characters')
1583
+ if not dashboard_id:
1584
+ raise InvalidUrnError("DashboardUrn dashboard_id cannot be empty")
1585
+ if UrnEncoder.contains_reserved_char(dashboard_id):
1586
+ raise InvalidUrnError(f'DashboardUrn dashboard_id contains reserved characters')
1350
1587
 
1351
- super().__init__(self.ENTITY_TYPE, [name, orchestrator, env])
1588
+ super().__init__(self.ENTITY_TYPE, [dashboard_tool, dashboard_id])
1352
1589
 
1353
1590
  @classmethod
1354
- def _parse_ids(cls, entity_ids: List[str]) -> "DataProcessUrn":
1591
+ def _parse_ids(cls, entity_ids: List[str]) -> "DashboardUrn":
1355
1592
  if len(entity_ids) != cls._URN_PARTS:
1356
- raise InvalidUrnError(f"DataProcessUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1357
- return cls(name=entity_ids[0], orchestrator=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
1593
+ raise InvalidUrnError(f"DashboardUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1594
+ return cls(dashboard_tool=entity_ids[0], dashboard_id=entity_ids[1], _allow_coercion=False)
1358
1595
 
1359
1596
  @classmethod
1360
- def underlying_key_aspect_type(cls) -> Type["DataProcessKeyClass"]:
1361
- from datahub.metadata.schema_classes import DataProcessKeyClass
1597
+ def underlying_key_aspect_type(cls) -> Type["DashboardKeyClass"]:
1598
+ from datahub.metadata.schema_classes import DashboardKeyClass
1362
1599
 
1363
- return DataProcessKeyClass
1600
+ return DashboardKeyClass
1364
1601
 
1365
- def to_key_aspect(self) -> "DataProcessKeyClass":
1366
- from datahub.metadata.schema_classes import DataProcessKeyClass
1602
+ def to_key_aspect(self) -> "DashboardKeyClass":
1603
+ from datahub.metadata.schema_classes import DashboardKeyClass
1367
1604
 
1368
- return DataProcessKeyClass(name=self.name, orchestrator=self.orchestrator, origin=self.env)
1605
+ return DashboardKeyClass(dashboardTool=self.dashboard_tool, dashboardId=self.dashboard_id)
1369
1606
 
1370
1607
  @classmethod
1371
- def from_key_aspect(cls, key_aspect: "DataProcessKeyClass") -> "DataProcessUrn":
1372
- return cls(name=key_aspect.name, orchestrator=key_aspect.orchestrator, env=key_aspect.origin)
1608
+ def from_key_aspect(cls, key_aspect: "DashboardKeyClass") -> "DashboardUrn":
1609
+ return cls(dashboard_tool=key_aspect.dashboardTool, dashboard_id=key_aspect.dashboardId)
1373
1610
 
1374
1611
  @property
1375
- def name(self) -> str:
1612
+ def dashboard_tool(self) -> str:
1376
1613
  return self.entity_ids[0]
1377
1614
 
1378
1615
  @property
1379
- def orchestrator(self) -> str:
1616
+ def dashboard_id(self) -> str:
1380
1617
  return self.entity_ids[1]
1381
1618
 
1382
- @property
1383
- def env(self) -> str:
1384
- return self.entity_ids[2]
1385
-
1386
1619
  if TYPE_CHECKING:
1387
- from datahub.metadata.schema_classes import DataHubViewKeyClass
1620
+ from datahub.metadata.schema_classes import TelemetryKeyClass
1388
1621
 
1389
- class DataHubViewUrn(_SpecificUrn):
1390
- ENTITY_TYPE: ClassVar[str] = "dataHubView"
1622
+ class TelemetryUrn(_SpecificUrn):
1623
+ ENTITY_TYPE: ClassVar[str] = "telemetry"
1391
1624
  _URN_PARTS: ClassVar[int] = 1
1392
1625
 
1393
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1626
+ def __init__(self, name: Union["TelemetryUrn", str], *, _allow_coercion: bool = True) -> None:
1394
1627
  if _allow_coercion:
1395
1628
  # Field coercion logic (if any is required).
1396
- id = UrnEncoder.encode_string(id)
1629
+ if isinstance(name, str):
1630
+ if name.startswith('urn:li:'):
1631
+ try:
1632
+ name = TelemetryUrn.from_string(name)
1633
+ except InvalidUrnError:
1634
+ raise InvalidUrnError(f'Expecting a TelemetryUrn but got {name}')
1635
+ else:
1636
+ name = UrnEncoder.encode_string(name)
1397
1637
 
1398
1638
  # Validation logic.
1399
- if not id:
1400
- raise InvalidUrnError("DataHubViewUrn id cannot be empty")
1401
- if UrnEncoder.contains_reserved_char(id):
1402
- raise InvalidUrnError(f'DataHubViewUrn id contains reserved characters')
1639
+ if not name:
1640
+ raise InvalidUrnError("TelemetryUrn name cannot be empty")
1641
+ if isinstance(name, TelemetryUrn):
1642
+ name = name.name
1643
+ elif isinstance(name, Urn):
1644
+ raise InvalidUrnError(f'Expecting a TelemetryUrn but got {name}')
1645
+ if UrnEncoder.contains_reserved_char(name):
1646
+ raise InvalidUrnError(f'TelemetryUrn name contains reserved characters')
1403
1647
 
1404
- super().__init__(self.ENTITY_TYPE, [id])
1648
+ super().__init__(self.ENTITY_TYPE, [name])
1405
1649
 
1406
1650
  @classmethod
1407
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubViewUrn":
1651
+ def _parse_ids(cls, entity_ids: List[str]) -> "TelemetryUrn":
1408
1652
  if len(entity_ids) != cls._URN_PARTS:
1409
- raise InvalidUrnError(f"DataHubViewUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1410
- return cls(id=entity_ids[0], _allow_coercion=False)
1653
+ raise InvalidUrnError(f"TelemetryUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1654
+ return cls(name=entity_ids[0], _allow_coercion=False)
1411
1655
 
1412
1656
  @classmethod
1413
- def underlying_key_aspect_type(cls) -> Type["DataHubViewKeyClass"]:
1414
- from datahub.metadata.schema_classes import DataHubViewKeyClass
1657
+ def underlying_key_aspect_type(cls) -> Type["TelemetryKeyClass"]:
1658
+ from datahub.metadata.schema_classes import TelemetryKeyClass
1415
1659
 
1416
- return DataHubViewKeyClass
1660
+ return TelemetryKeyClass
1417
1661
 
1418
- def to_key_aspect(self) -> "DataHubViewKeyClass":
1419
- from datahub.metadata.schema_classes import DataHubViewKeyClass
1662
+ def to_key_aspect(self) -> "TelemetryKeyClass":
1663
+ from datahub.metadata.schema_classes import TelemetryKeyClass
1420
1664
 
1421
- return DataHubViewKeyClass(id=self.id)
1665
+ return TelemetryKeyClass(name=self.name)
1422
1666
 
1423
1667
  @classmethod
1424
- def from_key_aspect(cls, key_aspect: "DataHubViewKeyClass") -> "DataHubViewUrn":
1425
- return cls(id=key_aspect.id)
1668
+ def from_key_aspect(cls, key_aspect: "TelemetryKeyClass") -> "TelemetryUrn":
1669
+ return cls(name=key_aspect.name)
1426
1670
 
1427
1671
  @property
1428
- def id(self) -> str:
1672
+ def name(self) -> str:
1429
1673
  return self.entity_ids[0]
1430
1674
 
1431
1675
  if TYPE_CHECKING:
1432
- from datahub.metadata.schema_classes import DataPlatformInstanceKeyClass
1676
+ from datahub.metadata.schema_classes import MLFeatureKeyClass
1433
1677
 
1434
- class DataPlatformInstanceUrn(_SpecificUrn):
1435
- ENTITY_TYPE: ClassVar[str] = "dataPlatformInstance"
1678
+ class MlFeatureUrn(_SpecificUrn):
1679
+ ENTITY_TYPE: ClassVar[str] = "mlFeature"
1436
1680
  _URN_PARTS: ClassVar[int] = 2
1437
1681
 
1438
- def __init__(self, platform: Union["DataPlatformUrn", str], instance: str, *, _allow_coercion: bool = True) -> None:
1682
+ def __init__(self, feature_namespace: str, name: str, *, _allow_coercion: bool = True) -> None:
1439
1683
  if _allow_coercion:
1440
1684
  # Field coercion logic (if any is required).
1441
- platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()
1442
- instance = UrnEncoder.encode_string(instance)
1685
+ feature_namespace = UrnEncoder.encode_string(feature_namespace)
1686
+ name = UrnEncoder.encode_string(name)
1443
1687
 
1444
1688
  # Validation logic.
1445
- if not platform:
1446
- raise InvalidUrnError("DataPlatformInstanceUrn platform cannot be empty")
1447
- platform = str(platform)
1448
- assert DataPlatformUrn.from_string(platform)
1449
- if not instance:
1450
- raise InvalidUrnError("DataPlatformInstanceUrn instance cannot be empty")
1451
- if UrnEncoder.contains_reserved_char(instance):
1452
- raise InvalidUrnError(f'DataPlatformInstanceUrn instance contains reserved characters')
1689
+ if not feature_namespace:
1690
+ raise InvalidUrnError("MlFeatureUrn feature_namespace cannot be empty")
1691
+ if UrnEncoder.contains_reserved_char(feature_namespace):
1692
+ raise InvalidUrnError(f'MlFeatureUrn feature_namespace contains reserved characters')
1693
+ if not name:
1694
+ raise InvalidUrnError("MlFeatureUrn name cannot be empty")
1695
+ if UrnEncoder.contains_reserved_char(name):
1696
+ raise InvalidUrnError(f'MlFeatureUrn name contains reserved characters')
1453
1697
 
1454
- super().__init__(self.ENTITY_TYPE, [platform, instance])
1698
+ super().__init__(self.ENTITY_TYPE, [feature_namespace, name])
1455
1699
 
1456
1700
  @classmethod
1457
- def _parse_ids(cls, entity_ids: List[str]) -> "DataPlatformInstanceUrn":
1701
+ def _parse_ids(cls, entity_ids: List[str]) -> "MlFeatureUrn":
1458
1702
  if len(entity_ids) != cls._URN_PARTS:
1459
- raise InvalidUrnError(f"DataPlatformInstanceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1460
- return cls(platform=entity_ids[0], instance=entity_ids[1], _allow_coercion=False)
1703
+ raise InvalidUrnError(f"MlFeatureUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1704
+ return cls(feature_namespace=entity_ids[0], name=entity_ids[1], _allow_coercion=False)
1461
1705
 
1462
1706
  @classmethod
1463
- def underlying_key_aspect_type(cls) -> Type["DataPlatformInstanceKeyClass"]:
1464
- from datahub.metadata.schema_classes import DataPlatformInstanceKeyClass
1707
+ def underlying_key_aspect_type(cls) -> Type["MLFeatureKeyClass"]:
1708
+ from datahub.metadata.schema_classes import MLFeatureKeyClass
1465
1709
 
1466
- return DataPlatformInstanceKeyClass
1710
+ return MLFeatureKeyClass
1467
1711
 
1468
- def to_key_aspect(self) -> "DataPlatformInstanceKeyClass":
1469
- from datahub.metadata.schema_classes import DataPlatformInstanceKeyClass
1712
+ def to_key_aspect(self) -> "MLFeatureKeyClass":
1713
+ from datahub.metadata.schema_classes import MLFeatureKeyClass
1470
1714
 
1471
- return DataPlatformInstanceKeyClass(platform=self.platform, instance=self.instance)
1715
+ return MLFeatureKeyClass(featureNamespace=self.feature_namespace, name=self.name)
1472
1716
 
1473
1717
  @classmethod
1474
- def from_key_aspect(cls, key_aspect: "DataPlatformInstanceKeyClass") -> "DataPlatformInstanceUrn":
1475
- return cls(platform=key_aspect.platform, instance=key_aspect.instance)
1718
+ def from_key_aspect(cls, key_aspect: "MLFeatureKeyClass") -> "MlFeatureUrn":
1719
+ return cls(feature_namespace=key_aspect.featureNamespace, name=key_aspect.name)
1476
1720
 
1477
1721
  @property
1478
- def platform(self) -> str:
1722
+ def feature_namespace(self) -> str:
1479
1723
  return self.entity_ids[0]
1480
1724
 
1481
1725
  @property
1482
- def instance(self) -> str:
1726
+ def name(self) -> str:
1483
1727
  return self.entity_ids[1]
1484
1728
 
1485
1729
  if TYPE_CHECKING:
1486
- from datahub.metadata.schema_classes import RoleKeyClass
1730
+ from datahub.metadata.schema_classes import DataHubStepStateKeyClass
1487
1731
 
1488
- class RoleUrn(_SpecificUrn):
1489
- ENTITY_TYPE: ClassVar[str] = "role"
1732
+ class DataHubStepStateUrn(_SpecificUrn):
1733
+ ENTITY_TYPE: ClassVar[str] = "dataHubStepState"
1490
1734
  _URN_PARTS: ClassVar[int] = 1
1491
1735
 
1492
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1736
+ def __init__(self, id: Union["DataHubStepStateUrn", str], *, _allow_coercion: bool = True) -> None:
1493
1737
  if _allow_coercion:
1494
1738
  # Field coercion logic (if any is required).
1495
- id = UrnEncoder.encode_string(id)
1739
+ if isinstance(id, str):
1740
+ if id.startswith('urn:li:'):
1741
+ try:
1742
+ id = DataHubStepStateUrn.from_string(id)
1743
+ except InvalidUrnError:
1744
+ raise InvalidUrnError(f'Expecting a DataHubStepStateUrn but got {id}')
1745
+ else:
1746
+ id = UrnEncoder.encode_string(id)
1496
1747
 
1497
1748
  # Validation logic.
1498
1749
  if not id:
1499
- raise InvalidUrnError("RoleUrn id cannot be empty")
1750
+ raise InvalidUrnError("DataHubStepStateUrn id cannot be empty")
1751
+ if isinstance(id, DataHubStepStateUrn):
1752
+ id = id.id
1753
+ elif isinstance(id, Urn):
1754
+ raise InvalidUrnError(f'Expecting a DataHubStepStateUrn but got {id}')
1500
1755
  if UrnEncoder.contains_reserved_char(id):
1501
- raise InvalidUrnError(f'RoleUrn id contains reserved characters')
1756
+ raise InvalidUrnError(f'DataHubStepStateUrn id contains reserved characters')
1502
1757
 
1503
1758
  super().__init__(self.ENTITY_TYPE, [id])
1504
1759
 
1505
1760
  @classmethod
1506
- def _parse_ids(cls, entity_ids: List[str]) -> "RoleUrn":
1761
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubStepStateUrn":
1507
1762
  if len(entity_ids) != cls._URN_PARTS:
1508
- raise InvalidUrnError(f"RoleUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1763
+ raise InvalidUrnError(f"DataHubStepStateUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1509
1764
  return cls(id=entity_ids[0], _allow_coercion=False)
1510
1765
 
1511
1766
  @classmethod
1512
- def underlying_key_aspect_type(cls) -> Type["RoleKeyClass"]:
1513
- from datahub.metadata.schema_classes import RoleKeyClass
1767
+ def underlying_key_aspect_type(cls) -> Type["DataHubStepStateKeyClass"]:
1768
+ from datahub.metadata.schema_classes import DataHubStepStateKeyClass
1514
1769
 
1515
- return RoleKeyClass
1770
+ return DataHubStepStateKeyClass
1516
1771
 
1517
- def to_key_aspect(self) -> "RoleKeyClass":
1518
- from datahub.metadata.schema_classes import RoleKeyClass
1772
+ def to_key_aspect(self) -> "DataHubStepStateKeyClass":
1773
+ from datahub.metadata.schema_classes import DataHubStepStateKeyClass
1519
1774
 
1520
- return RoleKeyClass(id=self.id)
1775
+ return DataHubStepStateKeyClass(id=self.id)
1521
1776
 
1522
1777
  @classmethod
1523
- def from_key_aspect(cls, key_aspect: "RoleKeyClass") -> "RoleUrn":
1778
+ def from_key_aspect(cls, key_aspect: "DataHubStepStateKeyClass") -> "DataHubStepStateUrn":
1524
1779
  return cls(id=key_aspect.id)
1525
1780
 
1526
1781
  @property
@@ -1528,44 +1783,55 @@ class RoleUrn(_SpecificUrn):
1528
1783
  return self.entity_ids[0]
1529
1784
 
1530
1785
  if TYPE_CHECKING:
1531
- from datahub.metadata.schema_classes import FormKeyClass
1786
+ from datahub.metadata.schema_classes import DataHubViewKeyClass
1532
1787
 
1533
- class FormUrn(_SpecificUrn):
1534
- ENTITY_TYPE: ClassVar[str] = "form"
1788
+ class DataHubViewUrn(_SpecificUrn):
1789
+ ENTITY_TYPE: ClassVar[str] = "dataHubView"
1535
1790
  _URN_PARTS: ClassVar[int] = 1
1536
1791
 
1537
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1792
+ def __init__(self, id: Union["DataHubViewUrn", str], *, _allow_coercion: bool = True) -> None:
1538
1793
  if _allow_coercion:
1539
1794
  # Field coercion logic (if any is required).
1540
- id = UrnEncoder.encode_string(id)
1795
+ if isinstance(id, str):
1796
+ if id.startswith('urn:li:'):
1797
+ try:
1798
+ id = DataHubViewUrn.from_string(id)
1799
+ except InvalidUrnError:
1800
+ raise InvalidUrnError(f'Expecting a DataHubViewUrn but got {id}')
1801
+ else:
1802
+ id = UrnEncoder.encode_string(id)
1541
1803
 
1542
1804
  # Validation logic.
1543
1805
  if not id:
1544
- raise InvalidUrnError("FormUrn id cannot be empty")
1806
+ raise InvalidUrnError("DataHubViewUrn id cannot be empty")
1807
+ if isinstance(id, DataHubViewUrn):
1808
+ id = id.id
1809
+ elif isinstance(id, Urn):
1810
+ raise InvalidUrnError(f'Expecting a DataHubViewUrn but got {id}')
1545
1811
  if UrnEncoder.contains_reserved_char(id):
1546
- raise InvalidUrnError(f'FormUrn id contains reserved characters')
1812
+ raise InvalidUrnError(f'DataHubViewUrn id contains reserved characters')
1547
1813
 
1548
1814
  super().__init__(self.ENTITY_TYPE, [id])
1549
1815
 
1550
1816
  @classmethod
1551
- def _parse_ids(cls, entity_ids: List[str]) -> "FormUrn":
1817
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubViewUrn":
1552
1818
  if len(entity_ids) != cls._URN_PARTS:
1553
- raise InvalidUrnError(f"FormUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1819
+ raise InvalidUrnError(f"DataHubViewUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1554
1820
  return cls(id=entity_ids[0], _allow_coercion=False)
1555
1821
 
1556
1822
  @classmethod
1557
- def underlying_key_aspect_type(cls) -> Type["FormKeyClass"]:
1558
- from datahub.metadata.schema_classes import FormKeyClass
1823
+ def underlying_key_aspect_type(cls) -> Type["DataHubViewKeyClass"]:
1824
+ from datahub.metadata.schema_classes import DataHubViewKeyClass
1559
1825
 
1560
- return FormKeyClass
1826
+ return DataHubViewKeyClass
1561
1827
 
1562
- def to_key_aspect(self) -> "FormKeyClass":
1563
- from datahub.metadata.schema_classes import FormKeyClass
1828
+ def to_key_aspect(self) -> "DataHubViewKeyClass":
1829
+ from datahub.metadata.schema_classes import DataHubViewKeyClass
1564
1830
 
1565
- return FormKeyClass(id=self.id)
1831
+ return DataHubViewKeyClass(id=self.id)
1566
1832
 
1567
1833
  @classmethod
1568
- def from_key_aspect(cls, key_aspect: "FormKeyClass") -> "FormUrn":
1834
+ def from_key_aspect(cls, key_aspect: "DataHubViewKeyClass") -> "DataHubViewUrn":
1569
1835
  return cls(id=key_aspect.id)
1570
1836
 
1571
1837
  @property
@@ -1573,478 +1839,490 @@ class FormUrn(_SpecificUrn):
1573
1839
  return self.entity_ids[0]
1574
1840
 
1575
1841
  if TYPE_CHECKING:
1576
- from datahub.metadata.schema_classes import MLFeatureKeyClass
1842
+ from datahub.metadata.schema_classes import DataHubRetentionKeyClass
1577
1843
 
1578
- class MlFeatureUrn(_SpecificUrn):
1579
- ENTITY_TYPE: ClassVar[str] = "mlFeature"
1844
+ class DataHubRetentionUrn(_SpecificUrn):
1845
+ ENTITY_TYPE: ClassVar[str] = "dataHubRetention"
1580
1846
  _URN_PARTS: ClassVar[int] = 2
1581
1847
 
1582
- def __init__(self, feature_namespace: str, name: str, *, _allow_coercion: bool = True) -> None:
1848
+ def __init__(self, entity_name: str, aspect_name: str, *, _allow_coercion: bool = True) -> None:
1583
1849
  if _allow_coercion:
1584
1850
  # Field coercion logic (if any is required).
1585
- feature_namespace = UrnEncoder.encode_string(feature_namespace)
1586
- name = UrnEncoder.encode_string(name)
1851
+ entity_name = UrnEncoder.encode_string(entity_name)
1852
+ aspect_name = UrnEncoder.encode_string(aspect_name)
1587
1853
 
1588
1854
  # Validation logic.
1589
- if not feature_namespace:
1590
- raise InvalidUrnError("MlFeatureUrn feature_namespace cannot be empty")
1591
- if UrnEncoder.contains_reserved_char(feature_namespace):
1592
- raise InvalidUrnError(f'MlFeatureUrn feature_namespace contains reserved characters')
1593
- if not name:
1594
- raise InvalidUrnError("MlFeatureUrn name cannot be empty")
1595
- if UrnEncoder.contains_reserved_char(name):
1596
- raise InvalidUrnError(f'MlFeatureUrn name contains reserved characters')
1855
+ if not entity_name:
1856
+ raise InvalidUrnError("DataHubRetentionUrn entity_name cannot be empty")
1857
+ if UrnEncoder.contains_reserved_char(entity_name):
1858
+ raise InvalidUrnError(f'DataHubRetentionUrn entity_name contains reserved characters')
1859
+ if not aspect_name:
1860
+ raise InvalidUrnError("DataHubRetentionUrn aspect_name cannot be empty")
1861
+ if UrnEncoder.contains_reserved_char(aspect_name):
1862
+ raise InvalidUrnError(f'DataHubRetentionUrn aspect_name contains reserved characters')
1597
1863
 
1598
- super().__init__(self.ENTITY_TYPE, [feature_namespace, name])
1864
+ super().__init__(self.ENTITY_TYPE, [entity_name, aspect_name])
1599
1865
 
1600
1866
  @classmethod
1601
- def _parse_ids(cls, entity_ids: List[str]) -> "MlFeatureUrn":
1867
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubRetentionUrn":
1602
1868
  if len(entity_ids) != cls._URN_PARTS:
1603
- raise InvalidUrnError(f"MlFeatureUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1604
- return cls(feature_namespace=entity_ids[0], name=entity_ids[1], _allow_coercion=False)
1869
+ raise InvalidUrnError(f"DataHubRetentionUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1870
+ return cls(entity_name=entity_ids[0], aspect_name=entity_ids[1], _allow_coercion=False)
1605
1871
 
1606
1872
  @classmethod
1607
- def underlying_key_aspect_type(cls) -> Type["MLFeatureKeyClass"]:
1608
- from datahub.metadata.schema_classes import MLFeatureKeyClass
1873
+ def underlying_key_aspect_type(cls) -> Type["DataHubRetentionKeyClass"]:
1874
+ from datahub.metadata.schema_classes import DataHubRetentionKeyClass
1609
1875
 
1610
- return MLFeatureKeyClass
1876
+ return DataHubRetentionKeyClass
1611
1877
 
1612
- def to_key_aspect(self) -> "MLFeatureKeyClass":
1613
- from datahub.metadata.schema_classes import MLFeatureKeyClass
1878
+ def to_key_aspect(self) -> "DataHubRetentionKeyClass":
1879
+ from datahub.metadata.schema_classes import DataHubRetentionKeyClass
1614
1880
 
1615
- return MLFeatureKeyClass(featureNamespace=self.feature_namespace, name=self.name)
1881
+ return DataHubRetentionKeyClass(entityName=self.entity_name, aspectName=self.aspect_name)
1616
1882
 
1617
1883
  @classmethod
1618
- def from_key_aspect(cls, key_aspect: "MLFeatureKeyClass") -> "MlFeatureUrn":
1619
- return cls(feature_namespace=key_aspect.featureNamespace, name=key_aspect.name)
1884
+ def from_key_aspect(cls, key_aspect: "DataHubRetentionKeyClass") -> "DataHubRetentionUrn":
1885
+ return cls(entity_name=key_aspect.entityName, aspect_name=key_aspect.aspectName)
1620
1886
 
1621
1887
  @property
1622
- def feature_namespace(self) -> str:
1888
+ def entity_name(self) -> str:
1623
1889
  return self.entity_ids[0]
1624
1890
 
1625
1891
  @property
1626
- def name(self) -> str:
1892
+ def aspect_name(self) -> str:
1627
1893
  return self.entity_ids[1]
1628
1894
 
1629
1895
  if TYPE_CHECKING:
1630
- from datahub.metadata.schema_classes import PostKeyClass
1896
+ from datahub.metadata.schema_classes import MLPrimaryKeyKeyClass
1631
1897
 
1632
- class PostUrn(_SpecificUrn):
1633
- ENTITY_TYPE: ClassVar[str] = "post"
1634
- _URN_PARTS: ClassVar[int] = 1
1898
+ class MlPrimaryKeyUrn(_SpecificUrn):
1899
+ ENTITY_TYPE: ClassVar[str] = "mlPrimaryKey"
1900
+ _URN_PARTS: ClassVar[int] = 2
1635
1901
 
1636
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
1902
+ def __init__(self, feature_namespace: str, name: str, *, _allow_coercion: bool = True) -> None:
1637
1903
  if _allow_coercion:
1638
1904
  # Field coercion logic (if any is required).
1639
- id = UrnEncoder.encode_string(id)
1905
+ feature_namespace = UrnEncoder.encode_string(feature_namespace)
1906
+ name = UrnEncoder.encode_string(name)
1640
1907
 
1641
1908
  # Validation logic.
1642
- if not id:
1643
- raise InvalidUrnError("PostUrn id cannot be empty")
1644
- if UrnEncoder.contains_reserved_char(id):
1645
- raise InvalidUrnError(f'PostUrn id contains reserved characters')
1909
+ if not feature_namespace:
1910
+ raise InvalidUrnError("MlPrimaryKeyUrn feature_namespace cannot be empty")
1911
+ if UrnEncoder.contains_reserved_char(feature_namespace):
1912
+ raise InvalidUrnError(f'MlPrimaryKeyUrn feature_namespace contains reserved characters')
1913
+ if not name:
1914
+ raise InvalidUrnError("MlPrimaryKeyUrn name cannot be empty")
1915
+ if UrnEncoder.contains_reserved_char(name):
1916
+ raise InvalidUrnError(f'MlPrimaryKeyUrn name contains reserved characters')
1646
1917
 
1647
- super().__init__(self.ENTITY_TYPE, [id])
1918
+ super().__init__(self.ENTITY_TYPE, [feature_namespace, name])
1648
1919
 
1649
1920
  @classmethod
1650
- def _parse_ids(cls, entity_ids: List[str]) -> "PostUrn":
1921
+ def _parse_ids(cls, entity_ids: List[str]) -> "MlPrimaryKeyUrn":
1651
1922
  if len(entity_ids) != cls._URN_PARTS:
1652
- raise InvalidUrnError(f"PostUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1653
- return cls(id=entity_ids[0], _allow_coercion=False)
1923
+ raise InvalidUrnError(f"MlPrimaryKeyUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1924
+ return cls(feature_namespace=entity_ids[0], name=entity_ids[1], _allow_coercion=False)
1654
1925
 
1655
1926
  @classmethod
1656
- def underlying_key_aspect_type(cls) -> Type["PostKeyClass"]:
1657
- from datahub.metadata.schema_classes import PostKeyClass
1927
+ def underlying_key_aspect_type(cls) -> Type["MLPrimaryKeyKeyClass"]:
1928
+ from datahub.metadata.schema_classes import MLPrimaryKeyKeyClass
1658
1929
 
1659
- return PostKeyClass
1930
+ return MLPrimaryKeyKeyClass
1660
1931
 
1661
- def to_key_aspect(self) -> "PostKeyClass":
1662
- from datahub.metadata.schema_classes import PostKeyClass
1932
+ def to_key_aspect(self) -> "MLPrimaryKeyKeyClass":
1933
+ from datahub.metadata.schema_classes import MLPrimaryKeyKeyClass
1663
1934
 
1664
- return PostKeyClass(id=self.id)
1935
+ return MLPrimaryKeyKeyClass(featureNamespace=self.feature_namespace, name=self.name)
1665
1936
 
1666
1937
  @classmethod
1667
- def from_key_aspect(cls, key_aspect: "PostKeyClass") -> "PostUrn":
1668
- return cls(id=key_aspect.id)
1938
+ def from_key_aspect(cls, key_aspect: "MLPrimaryKeyKeyClass") -> "MlPrimaryKeyUrn":
1939
+ return cls(feature_namespace=key_aspect.featureNamespace, name=key_aspect.name)
1669
1940
 
1670
1941
  @property
1671
- def id(self) -> str:
1942
+ def feature_namespace(self) -> str:
1672
1943
  return self.entity_ids[0]
1673
1944
 
1945
+ @property
1946
+ def name(self) -> str:
1947
+ return self.entity_ids[1]
1948
+
1674
1949
  if TYPE_CHECKING:
1675
- from datahub.metadata.schema_classes import DatasetKeyClass
1950
+ from datahub.metadata.schema_classes import DataFlowKeyClass
1676
1951
 
1677
- class DatasetUrn(_SpecificUrn):
1678
- ENTITY_TYPE: ClassVar[str] = "dataset"
1952
+ class DataFlowUrn(_SpecificUrn):
1953
+ ENTITY_TYPE: ClassVar[str] = "dataFlow"
1679
1954
  _URN_PARTS: ClassVar[int] = 3
1680
1955
 
1681
- def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
1956
+ def __init__(self, orchestrator: str, flow_id: str, cluster: str, *, _allow_coercion: bool = True) -> None:
1682
1957
  if _allow_coercion:
1683
1958
  # Field coercion logic (if any is required).
1684
- platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()
1685
- name = UrnEncoder.encode_string(name)
1686
- env = env.upper()
1687
- env = UrnEncoder.encode_string(env)
1959
+ orchestrator = UrnEncoder.encode_string(orchestrator)
1960
+ flow_id = UrnEncoder.encode_string(flow_id)
1961
+ cluster = UrnEncoder.encode_string(cluster)
1688
1962
 
1689
1963
  # Validation logic.
1690
- if not platform:
1691
- raise InvalidUrnError("DatasetUrn platform cannot be empty")
1692
- platform = str(platform)
1693
- assert DataPlatformUrn.from_string(platform)
1694
- if not name:
1695
- raise InvalidUrnError("DatasetUrn name cannot be empty")
1696
- if UrnEncoder.contains_reserved_char(name):
1697
- raise InvalidUrnError(f'DatasetUrn name contains reserved characters')
1698
- if not env:
1699
- raise InvalidUrnError("DatasetUrn env cannot be empty")
1700
- if UrnEncoder.contains_reserved_char(env):
1701
- raise InvalidUrnError(f'DatasetUrn env contains reserved characters')
1964
+ if not orchestrator:
1965
+ raise InvalidUrnError("DataFlowUrn orchestrator cannot be empty")
1966
+ if UrnEncoder.contains_reserved_char(orchestrator):
1967
+ raise InvalidUrnError(f'DataFlowUrn orchestrator contains reserved characters')
1968
+ if not flow_id:
1969
+ raise InvalidUrnError("DataFlowUrn flow_id cannot be empty")
1970
+ if UrnEncoder.contains_reserved_char(flow_id):
1971
+ raise InvalidUrnError(f'DataFlowUrn flow_id contains reserved characters')
1972
+ if not cluster:
1973
+ raise InvalidUrnError("DataFlowUrn cluster cannot be empty")
1974
+ if UrnEncoder.contains_reserved_char(cluster):
1975
+ raise InvalidUrnError(f'DataFlowUrn cluster contains reserved characters')
1702
1976
 
1703
- super().__init__(self.ENTITY_TYPE, [platform, name, env])
1977
+ super().__init__(self.ENTITY_TYPE, [orchestrator, flow_id, cluster])
1704
1978
 
1705
1979
  @classmethod
1706
- def _parse_ids(cls, entity_ids: List[str]) -> "DatasetUrn":
1980
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataFlowUrn":
1707
1981
  if len(entity_ids) != cls._URN_PARTS:
1708
- raise InvalidUrnError(f"DatasetUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1709
- return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
1982
+ raise InvalidUrnError(f"DataFlowUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1983
+ return cls(orchestrator=entity_ids[0], flow_id=entity_ids[1], cluster=entity_ids[2], _allow_coercion=False)
1710
1984
 
1711
1985
  @classmethod
1712
- def underlying_key_aspect_type(cls) -> Type["DatasetKeyClass"]:
1713
- from datahub.metadata.schema_classes import DatasetKeyClass
1986
+ def underlying_key_aspect_type(cls) -> Type["DataFlowKeyClass"]:
1987
+ from datahub.metadata.schema_classes import DataFlowKeyClass
1714
1988
 
1715
- return DatasetKeyClass
1989
+ return DataFlowKeyClass
1716
1990
 
1717
- def to_key_aspect(self) -> "DatasetKeyClass":
1718
- from datahub.metadata.schema_classes import DatasetKeyClass
1991
+ def to_key_aspect(self) -> "DataFlowKeyClass":
1992
+ from datahub.metadata.schema_classes import DataFlowKeyClass
1719
1993
 
1720
- return DatasetKeyClass(platform=self.platform, name=self.name, origin=self.env)
1994
+ return DataFlowKeyClass(orchestrator=self.orchestrator, flowId=self.flow_id, cluster=self.cluster)
1721
1995
 
1722
1996
  @classmethod
1723
- def from_key_aspect(cls, key_aspect: "DatasetKeyClass") -> "DatasetUrn":
1724
- return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
1997
+ def from_key_aspect(cls, key_aspect: "DataFlowKeyClass") -> "DataFlowUrn":
1998
+ return cls(orchestrator=key_aspect.orchestrator, flow_id=key_aspect.flowId, cluster=key_aspect.cluster)
1725
1999
 
1726
2000
  @classmethod
1727
2001
  def create_from_ids(
1728
2002
  cls,
1729
- platform_id: str,
1730
- table_name: str,
2003
+ orchestrator: str,
2004
+ flow_id: str,
1731
2005
  env: str,
1732
2006
  platform_instance: Optional[str] = None,
1733
- ) -> "DatasetUrn":
1734
- return DatasetUrn(
1735
- platform=platform_id,
1736
- name=f"{platform_instance}.{table_name}" if platform_instance else table_name,
1737
- env=env,
2007
+ ) -> "DataFlowUrn":
2008
+ return cls(
2009
+ orchestrator=orchestrator,
2010
+ flow_id=f"{platform_instance}.{flow_id}" if platform_instance else flow_id,
2011
+ cluster=env,
1738
2012
  )
1739
2013
 
1740
- from datahub.utilities.urns.field_paths import get_simple_field_path_from_v2_field_path as _get_simple_field_path_from_v2_field_path
1741
-
1742
- get_simple_field_path_from_v2_field_path = staticmethod(deprecated(reason='Use the function from the field_paths module instead')(_get_simple_field_path_from_v2_field_path))
1743
-
1744
- def get_data_platform_urn(self) -> "DataPlatformUrn":
1745
- return DataPlatformUrn.from_string(self.platform)
1746
-
1747
- @deprecated(reason="Use .name instead")
1748
- def get_dataset_name(self) -> str:
1749
- return self.name
2014
+ @deprecated(reason="Use .orchestrator instead")
2015
+ def get_orchestrator_name(self) -> str:
2016
+ return self.orchestrator
1750
2017
 
1751
- @deprecated(reason="Use .env instead")
2018
+ @deprecated(reason="Use .flow_id instead")
2019
+ def get_flow_id(self) -> str:
2020
+ return self.flow_id
2021
+
2022
+ @deprecated(reason="Use .cluster instead")
1752
2023
  def get_env(self) -> str:
1753
- return self.env
2024
+ return self.cluster
1754
2025
 
1755
2026
  @property
1756
- def platform(self) -> str:
2027
+ def orchestrator(self) -> str:
1757
2028
  return self.entity_ids[0]
1758
2029
 
1759
2030
  @property
1760
- def name(self) -> str:
2031
+ def flow_id(self) -> str:
1761
2032
  return self.entity_ids[1]
1762
2033
 
1763
2034
  @property
1764
- def env(self) -> str:
2035
+ def cluster(self) -> str:
1765
2036
  return self.entity_ids[2]
1766
2037
 
1767
2038
  if TYPE_CHECKING:
1768
- from datahub.metadata.schema_classes import MLModelKeyClass
2039
+ from datahub.metadata.schema_classes import CorpUserKeyClass
1769
2040
 
1770
- class MlModelUrn(_SpecificUrn):
1771
- ENTITY_TYPE: ClassVar[str] = "mlModel"
1772
- _URN_PARTS: ClassVar[int] = 3
2041
+ class CorpUserUrn(_SpecificUrn):
2042
+ ENTITY_TYPE: ClassVar[str] = "corpuser"
2043
+ _URN_PARTS: ClassVar[int] = 1
1773
2044
 
1774
- def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
2045
+ def __init__(self, username: Union["CorpUserUrn", str], *, _allow_coercion: bool = True) -> None:
1775
2046
  if _allow_coercion:
1776
2047
  # Field coercion logic (if any is required).
1777
- platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()
1778
- name = UrnEncoder.encode_string(name)
1779
- env = env.upper()
1780
- env = UrnEncoder.encode_string(env)
2048
+ if isinstance(username, str):
2049
+ if username.startswith('urn:li:'):
2050
+ try:
2051
+ username = CorpUserUrn.from_string(username)
2052
+ except InvalidUrnError:
2053
+ raise InvalidUrnError(f'Expecting a CorpUserUrn but got {username}')
2054
+ else:
2055
+ username = UrnEncoder.encode_string(username)
1781
2056
 
1782
2057
  # Validation logic.
1783
- if not platform:
1784
- raise InvalidUrnError("MlModelUrn platform cannot be empty")
1785
- platform = str(platform)
1786
- assert DataPlatformUrn.from_string(platform)
1787
- if not name:
1788
- raise InvalidUrnError("MlModelUrn name cannot be empty")
1789
- if UrnEncoder.contains_reserved_char(name):
1790
- raise InvalidUrnError(f'MlModelUrn name contains reserved characters')
1791
- if not env:
1792
- raise InvalidUrnError("MlModelUrn env cannot be empty")
1793
- if UrnEncoder.contains_reserved_char(env):
1794
- raise InvalidUrnError(f'MlModelUrn env contains reserved characters')
2058
+ if not username:
2059
+ raise InvalidUrnError("CorpUserUrn username cannot be empty")
2060
+ if isinstance(username, CorpUserUrn):
2061
+ username = username.username
2062
+ elif isinstance(username, Urn):
2063
+ raise InvalidUrnError(f'Expecting a CorpUserUrn but got {username}')
2064
+ if UrnEncoder.contains_reserved_char(username):
2065
+ raise InvalidUrnError(f'CorpUserUrn username contains reserved characters')
1795
2066
 
1796
- super().__init__(self.ENTITY_TYPE, [platform, name, env])
2067
+ super().__init__(self.ENTITY_TYPE, [username])
1797
2068
 
1798
2069
  @classmethod
1799
- def _parse_ids(cls, entity_ids: List[str]) -> "MlModelUrn":
2070
+ def _parse_ids(cls, entity_ids: List[str]) -> "CorpUserUrn":
1800
2071
  if len(entity_ids) != cls._URN_PARTS:
1801
- raise InvalidUrnError(f"MlModelUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1802
- return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
2072
+ raise InvalidUrnError(f"CorpUserUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2073
+ return cls(username=entity_ids[0], _allow_coercion=False)
1803
2074
 
1804
2075
  @classmethod
1805
- def underlying_key_aspect_type(cls) -> Type["MLModelKeyClass"]:
1806
- from datahub.metadata.schema_classes import MLModelKeyClass
2076
+ def underlying_key_aspect_type(cls) -> Type["CorpUserKeyClass"]:
2077
+ from datahub.metadata.schema_classes import CorpUserKeyClass
1807
2078
 
1808
- return MLModelKeyClass
2079
+ return CorpUserKeyClass
1809
2080
 
1810
- def to_key_aspect(self) -> "MLModelKeyClass":
1811
- from datahub.metadata.schema_classes import MLModelKeyClass
2081
+ def to_key_aspect(self) -> "CorpUserKeyClass":
2082
+ from datahub.metadata.schema_classes import CorpUserKeyClass
1812
2083
 
1813
- return MLModelKeyClass(platform=self.platform, name=self.name, origin=self.env)
2084
+ return CorpUserKeyClass(username=self.username)
1814
2085
 
1815
2086
  @classmethod
1816
- def from_key_aspect(cls, key_aspect: "MLModelKeyClass") -> "MlModelUrn":
1817
- return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
1818
-
1819
- @property
1820
- def platform(self) -> str:
1821
- return self.entity_ids[0]
2087
+ def from_key_aspect(cls, key_aspect: "CorpUserKeyClass") -> "CorpUserUrn":
2088
+ return cls(username=key_aspect.username)
1822
2089
 
1823
- @property
1824
- def name(self) -> str:
1825
- return self.entity_ids[1]
2090
+ @classmethod
2091
+ @deprecated(reason="Use the constructor instead")
2092
+ def create_from_id(cls, id: str) -> "CorpUserUrn":
2093
+ return cls(id)
1826
2094
 
1827
2095
  @property
1828
- def env(self) -> str:
1829
- return self.entity_ids[2]
2096
+ def username(self) -> str:
2097
+ return self.entity_ids[0]
1830
2098
 
1831
2099
  if TYPE_CHECKING:
1832
- from datahub.metadata.schema_classes import MLPrimaryKeyKeyClass
2100
+ from datahub.metadata.schema_classes import InviteTokenKeyClass
1833
2101
 
1834
- class MlPrimaryKeyUrn(_SpecificUrn):
1835
- ENTITY_TYPE: ClassVar[str] = "mlPrimaryKey"
1836
- _URN_PARTS: ClassVar[int] = 2
2102
+ class InviteTokenUrn(_SpecificUrn):
2103
+ ENTITY_TYPE: ClassVar[str] = "inviteToken"
2104
+ _URN_PARTS: ClassVar[int] = 1
1837
2105
 
1838
- def __init__(self, feature_namespace: str, name: str, *, _allow_coercion: bool = True) -> None:
2106
+ def __init__(self, id: Union["InviteTokenUrn", str], *, _allow_coercion: bool = True) -> None:
1839
2107
  if _allow_coercion:
1840
2108
  # Field coercion logic (if any is required).
1841
- feature_namespace = UrnEncoder.encode_string(feature_namespace)
1842
- name = UrnEncoder.encode_string(name)
2109
+ if isinstance(id, str):
2110
+ if id.startswith('urn:li:'):
2111
+ try:
2112
+ id = InviteTokenUrn.from_string(id)
2113
+ except InvalidUrnError:
2114
+ raise InvalidUrnError(f'Expecting a InviteTokenUrn but got {id}')
2115
+ else:
2116
+ id = UrnEncoder.encode_string(id)
1843
2117
 
1844
2118
  # Validation logic.
1845
- if not feature_namespace:
1846
- raise InvalidUrnError("MlPrimaryKeyUrn feature_namespace cannot be empty")
1847
- if UrnEncoder.contains_reserved_char(feature_namespace):
1848
- raise InvalidUrnError(f'MlPrimaryKeyUrn feature_namespace contains reserved characters')
1849
- if not name:
1850
- raise InvalidUrnError("MlPrimaryKeyUrn name cannot be empty")
1851
- if UrnEncoder.contains_reserved_char(name):
1852
- raise InvalidUrnError(f'MlPrimaryKeyUrn name contains reserved characters')
2119
+ if not id:
2120
+ raise InvalidUrnError("InviteTokenUrn id cannot be empty")
2121
+ if isinstance(id, InviteTokenUrn):
2122
+ id = id.id
2123
+ elif isinstance(id, Urn):
2124
+ raise InvalidUrnError(f'Expecting a InviteTokenUrn but got {id}')
2125
+ if UrnEncoder.contains_reserved_char(id):
2126
+ raise InvalidUrnError(f'InviteTokenUrn id contains reserved characters')
1853
2127
 
1854
- super().__init__(self.ENTITY_TYPE, [feature_namespace, name])
2128
+ super().__init__(self.ENTITY_TYPE, [id])
1855
2129
 
1856
2130
  @classmethod
1857
- def _parse_ids(cls, entity_ids: List[str]) -> "MlPrimaryKeyUrn":
2131
+ def _parse_ids(cls, entity_ids: List[str]) -> "InviteTokenUrn":
1858
2132
  if len(entity_ids) != cls._URN_PARTS:
1859
- raise InvalidUrnError(f"MlPrimaryKeyUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1860
- return cls(feature_namespace=entity_ids[0], name=entity_ids[1], _allow_coercion=False)
2133
+ raise InvalidUrnError(f"InviteTokenUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2134
+ return cls(id=entity_ids[0], _allow_coercion=False)
1861
2135
 
1862
2136
  @classmethod
1863
- def underlying_key_aspect_type(cls) -> Type["MLPrimaryKeyKeyClass"]:
1864
- from datahub.metadata.schema_classes import MLPrimaryKeyKeyClass
2137
+ def underlying_key_aspect_type(cls) -> Type["InviteTokenKeyClass"]:
2138
+ from datahub.metadata.schema_classes import InviteTokenKeyClass
1865
2139
 
1866
- return MLPrimaryKeyKeyClass
2140
+ return InviteTokenKeyClass
1867
2141
 
1868
- def to_key_aspect(self) -> "MLPrimaryKeyKeyClass":
1869
- from datahub.metadata.schema_classes import MLPrimaryKeyKeyClass
2142
+ def to_key_aspect(self) -> "InviteTokenKeyClass":
2143
+ from datahub.metadata.schema_classes import InviteTokenKeyClass
1870
2144
 
1871
- return MLPrimaryKeyKeyClass(featureNamespace=self.feature_namespace, name=self.name)
2145
+ return InviteTokenKeyClass(id=self.id)
1872
2146
 
1873
2147
  @classmethod
1874
- def from_key_aspect(cls, key_aspect: "MLPrimaryKeyKeyClass") -> "MlPrimaryKeyUrn":
1875
- return cls(feature_namespace=key_aspect.featureNamespace, name=key_aspect.name)
2148
+ def from_key_aspect(cls, key_aspect: "InviteTokenKeyClass") -> "InviteTokenUrn":
2149
+ return cls(id=key_aspect.id)
1876
2150
 
1877
2151
  @property
1878
- def feature_namespace(self) -> str:
2152
+ def id(self) -> str:
1879
2153
  return self.entity_ids[0]
1880
2154
 
1881
- @property
1882
- def name(self) -> str:
1883
- return self.entity_ids[1]
1884
-
1885
2155
  if TYPE_CHECKING:
1886
- from datahub.metadata.schema_classes import MLModelDeploymentKeyClass
2156
+ from datahub.metadata.schema_classes import DataPlatformKeyClass
1887
2157
 
1888
- class MlModelDeploymentUrn(_SpecificUrn):
1889
- ENTITY_TYPE: ClassVar[str] = "mlModelDeployment"
1890
- _URN_PARTS: ClassVar[int] = 3
2158
+ class DataPlatformUrn(_SpecificUrn):
2159
+ ENTITY_TYPE: ClassVar[str] = "dataPlatform"
2160
+ _URN_PARTS: ClassVar[int] = 1
1891
2161
 
1892
- def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
2162
+ def __init__(self, platform_name: Union["DataPlatformUrn", str], *, _allow_coercion: bool = True) -> None:
1893
2163
  if _allow_coercion:
1894
2164
  # Field coercion logic (if any is required).
1895
- platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()
1896
- name = UrnEncoder.encode_string(name)
1897
- env = env.upper()
1898
- env = UrnEncoder.encode_string(env)
2165
+ if isinstance(platform_name, str):
2166
+ if platform_name.startswith('urn:li:'):
2167
+ try:
2168
+ platform_name = DataPlatformUrn.from_string(platform_name)
2169
+ except InvalidUrnError:
2170
+ raise InvalidUrnError(f'Expecting a DataPlatformUrn but got {platform_name}')
2171
+ else:
2172
+ platform_name = UrnEncoder.encode_string(platform_name)
1899
2173
 
1900
2174
  # Validation logic.
1901
- if not platform:
1902
- raise InvalidUrnError("MlModelDeploymentUrn platform cannot be empty")
1903
- platform = str(platform)
1904
- assert DataPlatformUrn.from_string(platform)
1905
- if not name:
1906
- raise InvalidUrnError("MlModelDeploymentUrn name cannot be empty")
1907
- if UrnEncoder.contains_reserved_char(name):
1908
- raise InvalidUrnError(f'MlModelDeploymentUrn name contains reserved characters')
1909
- if not env:
1910
- raise InvalidUrnError("MlModelDeploymentUrn env cannot be empty")
1911
- if UrnEncoder.contains_reserved_char(env):
1912
- raise InvalidUrnError(f'MlModelDeploymentUrn env contains reserved characters')
2175
+ if not platform_name:
2176
+ raise InvalidUrnError("DataPlatformUrn platform_name cannot be empty")
2177
+ if isinstance(platform_name, DataPlatformUrn):
2178
+ platform_name = platform_name.platform_name
2179
+ elif isinstance(platform_name, Urn):
2180
+ raise InvalidUrnError(f'Expecting a DataPlatformUrn but got {platform_name}')
2181
+ if UrnEncoder.contains_reserved_char(platform_name):
2182
+ raise InvalidUrnError(f'DataPlatformUrn platform_name contains reserved characters')
1913
2183
 
1914
- super().__init__(self.ENTITY_TYPE, [platform, name, env])
2184
+ super().__init__(self.ENTITY_TYPE, [platform_name])
1915
2185
 
1916
2186
  @classmethod
1917
- def _parse_ids(cls, entity_ids: List[str]) -> "MlModelDeploymentUrn":
2187
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataPlatformUrn":
1918
2188
  if len(entity_ids) != cls._URN_PARTS:
1919
- raise InvalidUrnError(f"MlModelDeploymentUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1920
- return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
2189
+ raise InvalidUrnError(f"DataPlatformUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2190
+ return cls(platform_name=entity_ids[0], _allow_coercion=False)
1921
2191
 
1922
2192
  @classmethod
1923
- def underlying_key_aspect_type(cls) -> Type["MLModelDeploymentKeyClass"]:
1924
- from datahub.metadata.schema_classes import MLModelDeploymentKeyClass
2193
+ def underlying_key_aspect_type(cls) -> Type["DataPlatformKeyClass"]:
2194
+ from datahub.metadata.schema_classes import DataPlatformKeyClass
1925
2195
 
1926
- return MLModelDeploymentKeyClass
2196
+ return DataPlatformKeyClass
1927
2197
 
1928
- def to_key_aspect(self) -> "MLModelDeploymentKeyClass":
1929
- from datahub.metadata.schema_classes import MLModelDeploymentKeyClass
2198
+ def to_key_aspect(self) -> "DataPlatformKeyClass":
2199
+ from datahub.metadata.schema_classes import DataPlatformKeyClass
1930
2200
 
1931
- return MLModelDeploymentKeyClass(platform=self.platform, name=self.name, origin=self.env)
2201
+ return DataPlatformKeyClass(platformName=self.platform_name)
1932
2202
 
1933
2203
  @classmethod
1934
- def from_key_aspect(cls, key_aspect: "MLModelDeploymentKeyClass") -> "MlModelDeploymentUrn":
1935
- return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
1936
-
1937
- @property
1938
- def platform(self) -> str:
1939
- return self.entity_ids[0]
2204
+ def from_key_aspect(cls, key_aspect: "DataPlatformKeyClass") -> "DataPlatformUrn":
2205
+ return cls(platform_name=key_aspect.platformName)
1940
2206
 
1941
- @property
1942
- def name(self) -> str:
1943
- return self.entity_ids[1]
2207
+ @classmethod
2208
+ @deprecated(reason="Use the constructor instead")
2209
+ def create_from_id(cls, id: str) -> "DataPlatformUrn":
2210
+ return cls(id)
1944
2211
 
1945
2212
  @property
1946
- def env(self) -> str:
1947
- return self.entity_ids[2]
2213
+ def platform_name(self) -> str:
2214
+ return self.entity_ids[0]
1948
2215
 
1949
2216
  if TYPE_CHECKING:
1950
- from datahub.metadata.schema_classes import MLFeatureTableKeyClass
2217
+ from datahub.metadata.schema_classes import ContainerKeyClass
1951
2218
 
1952
- class MlFeatureTableUrn(_SpecificUrn):
1953
- ENTITY_TYPE: ClassVar[str] = "mlFeatureTable"
1954
- _URN_PARTS: ClassVar[int] = 2
2219
+ class ContainerUrn(_SpecificUrn):
2220
+ ENTITY_TYPE: ClassVar[str] = "container"
2221
+ _URN_PARTS: ClassVar[int] = 1
1955
2222
 
1956
- def __init__(self, platform: Union["DataPlatformUrn", str], name: str, *, _allow_coercion: bool = True) -> None:
2223
+ def __init__(self, guid: Union["ContainerUrn", str], *, _allow_coercion: bool = True) -> None:
1957
2224
  if _allow_coercion:
1958
2225
  # Field coercion logic (if any is required).
1959
- platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()
1960
- name = UrnEncoder.encode_string(name)
2226
+ if isinstance(guid, str):
2227
+ if guid.startswith('urn:li:'):
2228
+ try:
2229
+ guid = ContainerUrn.from_string(guid)
2230
+ except InvalidUrnError:
2231
+ raise InvalidUrnError(f'Expecting a ContainerUrn but got {guid}')
2232
+ else:
2233
+ guid = UrnEncoder.encode_string(guid)
1961
2234
 
1962
2235
  # Validation logic.
1963
- if not platform:
1964
- raise InvalidUrnError("MlFeatureTableUrn platform cannot be empty")
1965
- platform = str(platform)
1966
- assert DataPlatformUrn.from_string(platform)
1967
- if not name:
1968
- raise InvalidUrnError("MlFeatureTableUrn name cannot be empty")
1969
- if UrnEncoder.contains_reserved_char(name):
1970
- raise InvalidUrnError(f'MlFeatureTableUrn name contains reserved characters')
2236
+ if not guid:
2237
+ raise InvalidUrnError("ContainerUrn guid cannot be empty")
2238
+ if isinstance(guid, ContainerUrn):
2239
+ guid = guid.guid
2240
+ elif isinstance(guid, Urn):
2241
+ raise InvalidUrnError(f'Expecting a ContainerUrn but got {guid}')
2242
+ if UrnEncoder.contains_reserved_char(guid):
2243
+ raise InvalidUrnError(f'ContainerUrn guid contains reserved characters')
1971
2244
 
1972
- super().__init__(self.ENTITY_TYPE, [platform, name])
2245
+ super().__init__(self.ENTITY_TYPE, [guid])
1973
2246
 
1974
2247
  @classmethod
1975
- def _parse_ids(cls, entity_ids: List[str]) -> "MlFeatureTableUrn":
2248
+ def _parse_ids(cls, entity_ids: List[str]) -> "ContainerUrn":
1976
2249
  if len(entity_ids) != cls._URN_PARTS:
1977
- raise InvalidUrnError(f"MlFeatureTableUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
1978
- return cls(platform=entity_ids[0], name=entity_ids[1], _allow_coercion=False)
2250
+ raise InvalidUrnError(f"ContainerUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2251
+ return cls(guid=entity_ids[0], _allow_coercion=False)
1979
2252
 
1980
2253
  @classmethod
1981
- def underlying_key_aspect_type(cls) -> Type["MLFeatureTableKeyClass"]:
1982
- from datahub.metadata.schema_classes import MLFeatureTableKeyClass
2254
+ def underlying_key_aspect_type(cls) -> Type["ContainerKeyClass"]:
2255
+ from datahub.metadata.schema_classes import ContainerKeyClass
1983
2256
 
1984
- return MLFeatureTableKeyClass
2257
+ return ContainerKeyClass
1985
2258
 
1986
- def to_key_aspect(self) -> "MLFeatureTableKeyClass":
1987
- from datahub.metadata.schema_classes import MLFeatureTableKeyClass
2259
+ def to_key_aspect(self) -> "ContainerKeyClass":
2260
+ from datahub.metadata.schema_classes import ContainerKeyClass
1988
2261
 
1989
- return MLFeatureTableKeyClass(platform=self.platform, name=self.name)
2262
+ return ContainerKeyClass(guid=self.guid)
1990
2263
 
1991
2264
  @classmethod
1992
- def from_key_aspect(cls, key_aspect: "MLFeatureTableKeyClass") -> "MlFeatureTableUrn":
1993
- return cls(platform=key_aspect.platform, name=key_aspect.name)
2265
+ def from_key_aspect(cls, key_aspect: "ContainerKeyClass") -> "ContainerUrn":
2266
+ return cls(guid=key_aspect.guid)
1994
2267
 
1995
2268
  @property
1996
- def platform(self) -> str:
2269
+ def guid(self) -> str:
1997
2270
  return self.entity_ids[0]
1998
2271
 
1999
- @property
2000
- def name(self) -> str:
2001
- return self.entity_ids[1]
2002
-
2003
2272
  if TYPE_CHECKING:
2004
- from datahub.metadata.schema_classes import DataHubUpgradeKeyClass
2273
+ from datahub.metadata.schema_classes import DataPlatformInstanceKeyClass
2005
2274
 
2006
- class DataHubUpgradeUrn(_SpecificUrn):
2007
- ENTITY_TYPE: ClassVar[str] = "dataHubUpgrade"
2008
- _URN_PARTS: ClassVar[int] = 1
2275
+ class DataPlatformInstanceUrn(_SpecificUrn):
2276
+ ENTITY_TYPE: ClassVar[str] = "dataPlatformInstance"
2277
+ _URN_PARTS: ClassVar[int] = 2
2009
2278
 
2010
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2279
+ def __init__(self, platform: Union["DataPlatformUrn", str], instance: str, *, _allow_coercion: bool = True) -> None:
2011
2280
  if _allow_coercion:
2012
2281
  # Field coercion logic (if any is required).
2013
- id = UrnEncoder.encode_string(id)
2282
+ platform = DataPlatformUrn(platform).urn()
2283
+ instance = UrnEncoder.encode_string(instance)
2014
2284
 
2015
2285
  # Validation logic.
2016
- if not id:
2017
- raise InvalidUrnError("DataHubUpgradeUrn id cannot be empty")
2018
- if UrnEncoder.contains_reserved_char(id):
2019
- raise InvalidUrnError(f'DataHubUpgradeUrn id contains reserved characters')
2286
+ if not platform:
2287
+ raise InvalidUrnError("DataPlatformInstanceUrn platform cannot be empty")
2288
+ platform = str(platform) # convert urn type to str
2289
+ assert DataPlatformUrn.from_string(platform)
2290
+ if not instance:
2291
+ raise InvalidUrnError("DataPlatformInstanceUrn instance cannot be empty")
2292
+ if UrnEncoder.contains_reserved_char(instance):
2293
+ raise InvalidUrnError(f'DataPlatformInstanceUrn instance contains reserved characters')
2020
2294
 
2021
- super().__init__(self.ENTITY_TYPE, [id])
2295
+ super().__init__(self.ENTITY_TYPE, [platform, instance])
2022
2296
 
2023
2297
  @classmethod
2024
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubUpgradeUrn":
2298
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataPlatformInstanceUrn":
2025
2299
  if len(entity_ids) != cls._URN_PARTS:
2026
- raise InvalidUrnError(f"DataHubUpgradeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2027
- return cls(id=entity_ids[0], _allow_coercion=False)
2300
+ raise InvalidUrnError(f"DataPlatformInstanceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2301
+ return cls(platform=entity_ids[0], instance=entity_ids[1], _allow_coercion=False)
2028
2302
 
2029
2303
  @classmethod
2030
- def underlying_key_aspect_type(cls) -> Type["DataHubUpgradeKeyClass"]:
2031
- from datahub.metadata.schema_classes import DataHubUpgradeKeyClass
2304
+ def underlying_key_aspect_type(cls) -> Type["DataPlatformInstanceKeyClass"]:
2305
+ from datahub.metadata.schema_classes import DataPlatformInstanceKeyClass
2032
2306
 
2033
- return DataHubUpgradeKeyClass
2307
+ return DataPlatformInstanceKeyClass
2034
2308
 
2035
- def to_key_aspect(self) -> "DataHubUpgradeKeyClass":
2036
- from datahub.metadata.schema_classes import DataHubUpgradeKeyClass
2309
+ def to_key_aspect(self) -> "DataPlatformInstanceKeyClass":
2310
+ from datahub.metadata.schema_classes import DataPlatformInstanceKeyClass
2037
2311
 
2038
- return DataHubUpgradeKeyClass(id=self.id)
2312
+ return DataPlatformInstanceKeyClass(platform=self.platform, instance=self.instance)
2039
2313
 
2040
2314
  @classmethod
2041
- def from_key_aspect(cls, key_aspect: "DataHubUpgradeKeyClass") -> "DataHubUpgradeUrn":
2042
- return cls(id=key_aspect.id)
2315
+ def from_key_aspect(cls, key_aspect: "DataPlatformInstanceKeyClass") -> "DataPlatformInstanceUrn":
2316
+ return cls(platform=key_aspect.platform, instance=key_aspect.instance)
2043
2317
 
2044
2318
  @property
2045
- def id(self) -> str:
2319
+ def platform(self) -> str:
2046
2320
  return self.entity_ids[0]
2047
2321
 
2322
+ @property
2323
+ def instance(self) -> str:
2324
+ return self.entity_ids[1]
2325
+
2048
2326
  if TYPE_CHECKING:
2049
2327
  from datahub.metadata.schema_classes import GlossaryNodeKeyClass
2050
2328
 
@@ -2052,14 +2330,25 @@ class GlossaryNodeUrn(_SpecificUrn):
2052
2330
  ENTITY_TYPE: ClassVar[str] = "glossaryNode"
2053
2331
  _URN_PARTS: ClassVar[int] = 1
2054
2332
 
2055
- def __init__(self, name: str, *, _allow_coercion: bool = True) -> None:
2333
+ def __init__(self, name: Union["GlossaryNodeUrn", str], *, _allow_coercion: bool = True) -> None:
2056
2334
  if _allow_coercion:
2057
2335
  # Field coercion logic (if any is required).
2058
- name = UrnEncoder.encode_string(name)
2336
+ if isinstance(name, str):
2337
+ if name.startswith('urn:li:'):
2338
+ try:
2339
+ name = GlossaryNodeUrn.from_string(name)
2340
+ except InvalidUrnError:
2341
+ raise InvalidUrnError(f'Expecting a GlossaryNodeUrn but got {name}')
2342
+ else:
2343
+ name = UrnEncoder.encode_string(name)
2059
2344
 
2060
2345
  # Validation logic.
2061
2346
  if not name:
2062
2347
  raise InvalidUrnError("GlossaryNodeUrn name cannot be empty")
2348
+ if isinstance(name, GlossaryNodeUrn):
2349
+ name = name.name
2350
+ elif isinstance(name, Urn):
2351
+ raise InvalidUrnError(f'Expecting a GlossaryNodeUrn but got {name}')
2063
2352
  if UrnEncoder.contains_reserved_char(name):
2064
2353
  raise InvalidUrnError(f'GlossaryNodeUrn name contains reserved characters')
2065
2354
 
@@ -2091,44 +2380,118 @@ class GlossaryNodeUrn(_SpecificUrn):
2091
2380
  return self.entity_ids[0]
2092
2381
 
2093
2382
  if TYPE_CHECKING:
2094
- from datahub.metadata.schema_classes import DataHubPersonaKeyClass
2383
+ from datahub.metadata.schema_classes import MLModelKeyClass
2095
2384
 
2096
- class DataHubPersonaUrn(_SpecificUrn):
2097
- ENTITY_TYPE: ClassVar[str] = "dataHubPersona"
2385
+ class MlModelUrn(_SpecificUrn):
2386
+ ENTITY_TYPE: ClassVar[str] = "mlModel"
2387
+ _URN_PARTS: ClassVar[int] = 3
2388
+
2389
+ def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
2390
+ if _allow_coercion:
2391
+ # Field coercion logic (if any is required).
2392
+ platform = DataPlatformUrn(platform).urn()
2393
+ name = UrnEncoder.encode_string(name)
2394
+ env = env.upper()
2395
+
2396
+ # Validation logic.
2397
+ if not platform:
2398
+ raise InvalidUrnError("MlModelUrn platform cannot be empty")
2399
+ platform = str(platform) # convert urn type to str
2400
+ assert DataPlatformUrn.from_string(platform)
2401
+ if not name:
2402
+ raise InvalidUrnError("MlModelUrn name cannot be empty")
2403
+ if UrnEncoder.contains_reserved_char(name):
2404
+ raise InvalidUrnError(f'MlModelUrn name contains reserved characters')
2405
+ if not env:
2406
+ raise InvalidUrnError("MlModelUrn env cannot be empty")
2407
+ if UrnEncoder.contains_reserved_char(env):
2408
+ raise InvalidUrnError(f'MlModelUrn env contains reserved characters')
2409
+
2410
+ super().__init__(self.ENTITY_TYPE, [platform, name, env])
2411
+
2412
+ @classmethod
2413
+ def _parse_ids(cls, entity_ids: List[str]) -> "MlModelUrn":
2414
+ if len(entity_ids) != cls._URN_PARTS:
2415
+ raise InvalidUrnError(f"MlModelUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2416
+ return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
2417
+
2418
+ @classmethod
2419
+ def underlying_key_aspect_type(cls) -> Type["MLModelKeyClass"]:
2420
+ from datahub.metadata.schema_classes import MLModelKeyClass
2421
+
2422
+ return MLModelKeyClass
2423
+
2424
+ def to_key_aspect(self) -> "MLModelKeyClass":
2425
+ from datahub.metadata.schema_classes import MLModelKeyClass
2426
+
2427
+ return MLModelKeyClass(platform=self.platform, name=self.name, origin=self.env)
2428
+
2429
+ @classmethod
2430
+ def from_key_aspect(cls, key_aspect: "MLModelKeyClass") -> "MlModelUrn":
2431
+ return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
2432
+
2433
+ @property
2434
+ def platform(self) -> str:
2435
+ return self.entity_ids[0]
2436
+
2437
+ @property
2438
+ def name(self) -> str:
2439
+ return self.entity_ids[1]
2440
+
2441
+ @property
2442
+ def env(self) -> str:
2443
+ return self.entity_ids[2]
2444
+
2445
+ if TYPE_CHECKING:
2446
+ from datahub.metadata.schema_classes import TestKeyClass
2447
+
2448
+ class TestUrn(_SpecificUrn):
2449
+ ENTITY_TYPE: ClassVar[str] = "test"
2098
2450
  _URN_PARTS: ClassVar[int] = 1
2099
2451
 
2100
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2452
+ def __init__(self, id: Union["TestUrn", str], *, _allow_coercion: bool = True) -> None:
2101
2453
  if _allow_coercion:
2102
2454
  # Field coercion logic (if any is required).
2103
- id = UrnEncoder.encode_string(id)
2455
+ if isinstance(id, str):
2456
+ if id.startswith('urn:li:'):
2457
+ try:
2458
+ id = TestUrn.from_string(id)
2459
+ except InvalidUrnError:
2460
+ raise InvalidUrnError(f'Expecting a TestUrn but got {id}')
2461
+ else:
2462
+ id = UrnEncoder.encode_string(id)
2104
2463
 
2105
2464
  # Validation logic.
2106
2465
  if not id:
2107
- raise InvalidUrnError("DataHubPersonaUrn id cannot be empty")
2466
+ raise InvalidUrnError("TestUrn id cannot be empty")
2467
+ if isinstance(id, TestUrn):
2468
+ id = id.id
2469
+ elif isinstance(id, Urn):
2470
+ raise InvalidUrnError(f'Expecting a TestUrn but got {id}')
2108
2471
  if UrnEncoder.contains_reserved_char(id):
2109
- raise InvalidUrnError(f'DataHubPersonaUrn id contains reserved characters')
2472
+ raise InvalidUrnError(f'TestUrn id contains reserved characters')
2110
2473
 
2111
2474
  super().__init__(self.ENTITY_TYPE, [id])
2112
2475
 
2113
2476
  @classmethod
2114
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubPersonaUrn":
2477
+ def _parse_ids(cls, entity_ids: List[str]) -> "TestUrn":
2115
2478
  if len(entity_ids) != cls._URN_PARTS:
2116
- raise InvalidUrnError(f"DataHubPersonaUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2479
+ raise InvalidUrnError(f"TestUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2117
2480
  return cls(id=entity_ids[0], _allow_coercion=False)
2118
2481
 
2119
2482
  @classmethod
2120
- def underlying_key_aspect_type(cls) -> Type["DataHubPersonaKeyClass"]:
2121
- from datahub.metadata.schema_classes import DataHubPersonaKeyClass
2483
+ def underlying_key_aspect_type(cls) -> Type["TestKeyClass"]:
2484
+ from datahub.metadata.schema_classes import TestKeyClass
2122
2485
 
2123
- return DataHubPersonaKeyClass
2486
+ return TestKeyClass
2124
2487
 
2125
- def to_key_aspect(self) -> "DataHubPersonaKeyClass":
2126
- from datahub.metadata.schema_classes import DataHubPersonaKeyClass
2488
+ def to_key_aspect(self) -> "TestKeyClass":
2489
+ from datahub.metadata.schema_classes import TestKeyClass
2127
2490
 
2128
- return DataHubPersonaKeyClass(id=self.id)
2491
+ return TestKeyClass(id=self.id)
2129
2492
 
2130
2493
  @classmethod
2131
- def from_key_aspect(cls, key_aspect: "DataHubPersonaKeyClass") -> "DataHubPersonaUrn":
2494
+ def from_key_aspect(cls, key_aspect: "TestKeyClass") -> "TestUrn":
2132
2495
  return cls(id=key_aspect.id)
2133
2496
 
2134
2497
  @property
@@ -2136,44 +2499,55 @@ class DataHubPersonaUrn(_SpecificUrn):
2136
2499
  return self.entity_ids[0]
2137
2500
 
2138
2501
  if TYPE_CHECKING:
2139
- from datahub.metadata.schema_classes import DataHubSecretKeyClass
2502
+ from datahub.metadata.schema_classes import PostKeyClass
2140
2503
 
2141
- class DataHubSecretUrn(_SpecificUrn):
2142
- ENTITY_TYPE: ClassVar[str] = "dataHubSecret"
2504
+ class PostUrn(_SpecificUrn):
2505
+ ENTITY_TYPE: ClassVar[str] = "post"
2143
2506
  _URN_PARTS: ClassVar[int] = 1
2144
2507
 
2145
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2508
+ def __init__(self, id: Union["PostUrn", str], *, _allow_coercion: bool = True) -> None:
2146
2509
  if _allow_coercion:
2147
2510
  # Field coercion logic (if any is required).
2148
- id = UrnEncoder.encode_string(id)
2511
+ if isinstance(id, str):
2512
+ if id.startswith('urn:li:'):
2513
+ try:
2514
+ id = PostUrn.from_string(id)
2515
+ except InvalidUrnError:
2516
+ raise InvalidUrnError(f'Expecting a PostUrn but got {id}')
2517
+ else:
2518
+ id = UrnEncoder.encode_string(id)
2149
2519
 
2150
2520
  # Validation logic.
2151
2521
  if not id:
2152
- raise InvalidUrnError("DataHubSecretUrn id cannot be empty")
2522
+ raise InvalidUrnError("PostUrn id cannot be empty")
2523
+ if isinstance(id, PostUrn):
2524
+ id = id.id
2525
+ elif isinstance(id, Urn):
2526
+ raise InvalidUrnError(f'Expecting a PostUrn but got {id}')
2153
2527
  if UrnEncoder.contains_reserved_char(id):
2154
- raise InvalidUrnError(f'DataHubSecretUrn id contains reserved characters')
2528
+ raise InvalidUrnError(f'PostUrn id contains reserved characters')
2155
2529
 
2156
2530
  super().__init__(self.ENTITY_TYPE, [id])
2157
2531
 
2158
2532
  @classmethod
2159
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubSecretUrn":
2533
+ def _parse_ids(cls, entity_ids: List[str]) -> "PostUrn":
2160
2534
  if len(entity_ids) != cls._URN_PARTS:
2161
- raise InvalidUrnError(f"DataHubSecretUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2535
+ raise InvalidUrnError(f"PostUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2162
2536
  return cls(id=entity_ids[0], _allow_coercion=False)
2163
2537
 
2164
2538
  @classmethod
2165
- def underlying_key_aspect_type(cls) -> Type["DataHubSecretKeyClass"]:
2166
- from datahub.metadata.schema_classes import DataHubSecretKeyClass
2539
+ def underlying_key_aspect_type(cls) -> Type["PostKeyClass"]:
2540
+ from datahub.metadata.schema_classes import PostKeyClass
2167
2541
 
2168
- return DataHubSecretKeyClass
2542
+ return PostKeyClass
2169
2543
 
2170
- def to_key_aspect(self) -> "DataHubSecretKeyClass":
2171
- from datahub.metadata.schema_classes import DataHubSecretKeyClass
2544
+ def to_key_aspect(self) -> "PostKeyClass":
2545
+ from datahub.metadata.schema_classes import PostKeyClass
2172
2546
 
2173
- return DataHubSecretKeyClass(id=self.id)
2547
+ return PostKeyClass(id=self.id)
2174
2548
 
2175
2549
  @classmethod
2176
- def from_key_aspect(cls, key_aspect: "DataHubSecretKeyClass") -> "DataHubSecretUrn":
2550
+ def from_key_aspect(cls, key_aspect: "PostKeyClass") -> "PostUrn":
2177
2551
  return cls(id=key_aspect.id)
2178
2552
 
2179
2553
  @property
@@ -2181,44 +2555,111 @@ class DataHubSecretUrn(_SpecificUrn):
2181
2555
  return self.entity_ids[0]
2182
2556
 
2183
2557
  if TYPE_CHECKING:
2184
- from datahub.metadata.schema_classes import DataHubIngestionSourceKeyClass
2558
+ from datahub.metadata.schema_classes import GlossaryTermKeyClass
2185
2559
 
2186
- class DataHubIngestionSourceUrn(_SpecificUrn):
2187
- ENTITY_TYPE: ClassVar[str] = "dataHubIngestionSource"
2560
+ class GlossaryTermUrn(_SpecificUrn):
2561
+ ENTITY_TYPE: ClassVar[str] = "glossaryTerm"
2188
2562
  _URN_PARTS: ClassVar[int] = 1
2189
2563
 
2190
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2564
+ def __init__(self, name: Union["GlossaryTermUrn", str], *, _allow_coercion: bool = True) -> None:
2191
2565
  if _allow_coercion:
2192
2566
  # Field coercion logic (if any is required).
2193
- id = UrnEncoder.encode_string(id)
2567
+ if isinstance(name, str):
2568
+ if name.startswith('urn:li:'):
2569
+ try:
2570
+ name = GlossaryTermUrn.from_string(name)
2571
+ except InvalidUrnError:
2572
+ raise InvalidUrnError(f'Expecting a GlossaryTermUrn but got {name}')
2573
+ else:
2574
+ name = UrnEncoder.encode_string(name)
2575
+
2576
+ # Validation logic.
2577
+ if not name:
2578
+ raise InvalidUrnError("GlossaryTermUrn name cannot be empty")
2579
+ if isinstance(name, GlossaryTermUrn):
2580
+ name = name.name
2581
+ elif isinstance(name, Urn):
2582
+ raise InvalidUrnError(f'Expecting a GlossaryTermUrn but got {name}')
2583
+ if UrnEncoder.contains_reserved_char(name):
2584
+ raise InvalidUrnError(f'GlossaryTermUrn name contains reserved characters')
2585
+
2586
+ super().__init__(self.ENTITY_TYPE, [name])
2587
+
2588
+ @classmethod
2589
+ def _parse_ids(cls, entity_ids: List[str]) -> "GlossaryTermUrn":
2590
+ if len(entity_ids) != cls._URN_PARTS:
2591
+ raise InvalidUrnError(f"GlossaryTermUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2592
+ return cls(name=entity_ids[0], _allow_coercion=False)
2593
+
2594
+ @classmethod
2595
+ def underlying_key_aspect_type(cls) -> Type["GlossaryTermKeyClass"]:
2596
+ from datahub.metadata.schema_classes import GlossaryTermKeyClass
2597
+
2598
+ return GlossaryTermKeyClass
2599
+
2600
+ def to_key_aspect(self) -> "GlossaryTermKeyClass":
2601
+ from datahub.metadata.schema_classes import GlossaryTermKeyClass
2602
+
2603
+ return GlossaryTermKeyClass(name=self.name)
2604
+
2605
+ @classmethod
2606
+ def from_key_aspect(cls, key_aspect: "GlossaryTermKeyClass") -> "GlossaryTermUrn":
2607
+ return cls(name=key_aspect.name)
2608
+
2609
+ @property
2610
+ def name(self) -> str:
2611
+ return self.entity_ids[0]
2612
+
2613
+ if TYPE_CHECKING:
2614
+ from datahub.metadata.schema_classes import ERModelRelationshipKeyClass
2615
+
2616
+ class ErModelRelationshipUrn(_SpecificUrn):
2617
+ ENTITY_TYPE: ClassVar[str] = "erModelRelationship"
2618
+ _URN_PARTS: ClassVar[int] = 1
2619
+
2620
+ def __init__(self, id: Union["ErModelRelationshipUrn", str], *, _allow_coercion: bool = True) -> None:
2621
+ if _allow_coercion:
2622
+ # Field coercion logic (if any is required).
2623
+ if isinstance(id, str):
2624
+ if id.startswith('urn:li:'):
2625
+ try:
2626
+ id = ErModelRelationshipUrn.from_string(id)
2627
+ except InvalidUrnError:
2628
+ raise InvalidUrnError(f'Expecting a ErModelRelationshipUrn but got {id}')
2629
+ else:
2630
+ id = UrnEncoder.encode_string(id)
2194
2631
 
2195
2632
  # Validation logic.
2196
2633
  if not id:
2197
- raise InvalidUrnError("DataHubIngestionSourceUrn id cannot be empty")
2634
+ raise InvalidUrnError("ErModelRelationshipUrn id cannot be empty")
2635
+ if isinstance(id, ErModelRelationshipUrn):
2636
+ id = id.id
2637
+ elif isinstance(id, Urn):
2638
+ raise InvalidUrnError(f'Expecting a ErModelRelationshipUrn but got {id}')
2198
2639
  if UrnEncoder.contains_reserved_char(id):
2199
- raise InvalidUrnError(f'DataHubIngestionSourceUrn id contains reserved characters')
2640
+ raise InvalidUrnError(f'ErModelRelationshipUrn id contains reserved characters')
2200
2641
 
2201
2642
  super().__init__(self.ENTITY_TYPE, [id])
2202
2643
 
2203
2644
  @classmethod
2204
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubIngestionSourceUrn":
2645
+ def _parse_ids(cls, entity_ids: List[str]) -> "ErModelRelationshipUrn":
2205
2646
  if len(entity_ids) != cls._URN_PARTS:
2206
- raise InvalidUrnError(f"DataHubIngestionSourceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2647
+ raise InvalidUrnError(f"ErModelRelationshipUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2207
2648
  return cls(id=entity_ids[0], _allow_coercion=False)
2208
2649
 
2209
2650
  @classmethod
2210
- def underlying_key_aspect_type(cls) -> Type["DataHubIngestionSourceKeyClass"]:
2211
- from datahub.metadata.schema_classes import DataHubIngestionSourceKeyClass
2651
+ def underlying_key_aspect_type(cls) -> Type["ERModelRelationshipKeyClass"]:
2652
+ from datahub.metadata.schema_classes import ERModelRelationshipKeyClass
2212
2653
 
2213
- return DataHubIngestionSourceKeyClass
2654
+ return ERModelRelationshipKeyClass
2214
2655
 
2215
- def to_key_aspect(self) -> "DataHubIngestionSourceKeyClass":
2216
- from datahub.metadata.schema_classes import DataHubIngestionSourceKeyClass
2656
+ def to_key_aspect(self) -> "ERModelRelationshipKeyClass":
2657
+ from datahub.metadata.schema_classes import ERModelRelationshipKeyClass
2217
2658
 
2218
- return DataHubIngestionSourceKeyClass(id=self.id)
2659
+ return ERModelRelationshipKeyClass(id=self.id)
2219
2660
 
2220
2661
  @classmethod
2221
- def from_key_aspect(cls, key_aspect: "DataHubIngestionSourceKeyClass") -> "DataHubIngestionSourceUrn":
2662
+ def from_key_aspect(cls, key_aspect: "ERModelRelationshipKeyClass") -> "ErModelRelationshipUrn":
2222
2663
  return cls(id=key_aspect.id)
2223
2664
 
2224
2665
  @property
@@ -2226,152 +2667,167 @@ class DataHubIngestionSourceUrn(_SpecificUrn):
2226
2667
  return self.entity_ids[0]
2227
2668
 
2228
2669
  if TYPE_CHECKING:
2229
- from datahub.metadata.schema_classes import VersionSetKeyClass
2670
+ from datahub.metadata.schema_classes import ExecutionRequestKeyClass
2230
2671
 
2231
- class VersionSetUrn(_SpecificUrn):
2232
- ENTITY_TYPE: ClassVar[str] = "versionSet"
2233
- _URN_PARTS: ClassVar[int] = 2
2672
+ class DataHubExecutionRequestUrn(_SpecificUrn):
2673
+ ENTITY_TYPE: ClassVar[str] = "dataHubExecutionRequest"
2674
+ _URN_PARTS: ClassVar[int] = 1
2234
2675
 
2235
- def __init__(self, id: str, entity_type: str, *, _allow_coercion: bool = True) -> None:
2676
+ def __init__(self, id: Union["DataHubExecutionRequestUrn", str], *, _allow_coercion: bool = True) -> None:
2236
2677
  if _allow_coercion:
2237
2678
  # Field coercion logic (if any is required).
2238
- id = UrnEncoder.encode_string(id)
2239
- entity_type = UrnEncoder.encode_string(entity_type)
2679
+ if isinstance(id, str):
2680
+ if id.startswith('urn:li:'):
2681
+ try:
2682
+ id = DataHubExecutionRequestUrn.from_string(id)
2683
+ except InvalidUrnError:
2684
+ raise InvalidUrnError(f'Expecting a DataHubExecutionRequestUrn but got {id}')
2685
+ else:
2686
+ id = UrnEncoder.encode_string(id)
2240
2687
 
2241
2688
  # Validation logic.
2242
2689
  if not id:
2243
- raise InvalidUrnError("VersionSetUrn id cannot be empty")
2690
+ raise InvalidUrnError("DataHubExecutionRequestUrn id cannot be empty")
2691
+ if isinstance(id, DataHubExecutionRequestUrn):
2692
+ id = id.id
2693
+ elif isinstance(id, Urn):
2694
+ raise InvalidUrnError(f'Expecting a DataHubExecutionRequestUrn but got {id}')
2244
2695
  if UrnEncoder.contains_reserved_char(id):
2245
- raise InvalidUrnError(f'VersionSetUrn id contains reserved characters')
2246
- if not entity_type:
2247
- raise InvalidUrnError("VersionSetUrn entity_type cannot be empty")
2248
- if UrnEncoder.contains_reserved_char(entity_type):
2249
- raise InvalidUrnError(f'VersionSetUrn entity_type contains reserved characters')
2696
+ raise InvalidUrnError(f'DataHubExecutionRequestUrn id contains reserved characters')
2250
2697
 
2251
- super().__init__(self.ENTITY_TYPE, [id, entity_type])
2698
+ super().__init__(self.ENTITY_TYPE, [id])
2252
2699
 
2253
2700
  @classmethod
2254
- def _parse_ids(cls, entity_ids: List[str]) -> "VersionSetUrn":
2701
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubExecutionRequestUrn":
2255
2702
  if len(entity_ids) != cls._URN_PARTS:
2256
- raise InvalidUrnError(f"VersionSetUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2257
- return cls(id=entity_ids[0], entity_type=entity_ids[1], _allow_coercion=False)
2703
+ raise InvalidUrnError(f"DataHubExecutionRequestUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2704
+ return cls(id=entity_ids[0], _allow_coercion=False)
2258
2705
 
2259
2706
  @classmethod
2260
- def underlying_key_aspect_type(cls) -> Type["VersionSetKeyClass"]:
2261
- from datahub.metadata.schema_classes import VersionSetKeyClass
2707
+ def underlying_key_aspect_type(cls) -> Type["ExecutionRequestKeyClass"]:
2708
+ from datahub.metadata.schema_classes import ExecutionRequestKeyClass
2262
2709
 
2263
- return VersionSetKeyClass
2710
+ return ExecutionRequestKeyClass
2264
2711
 
2265
- def to_key_aspect(self) -> "VersionSetKeyClass":
2266
- from datahub.metadata.schema_classes import VersionSetKeyClass
2712
+ def to_key_aspect(self) -> "ExecutionRequestKeyClass":
2713
+ from datahub.metadata.schema_classes import ExecutionRequestKeyClass
2267
2714
 
2268
- return VersionSetKeyClass(id=self.id, entityType=self.entity_type)
2715
+ return ExecutionRequestKeyClass(id=self.id)
2269
2716
 
2270
2717
  @classmethod
2271
- def from_key_aspect(cls, key_aspect: "VersionSetKeyClass") -> "VersionSetUrn":
2272
- return cls(id=key_aspect.id, entity_type=key_aspect.entityType)
2718
+ def from_key_aspect(cls, key_aspect: "ExecutionRequestKeyClass") -> "DataHubExecutionRequestUrn":
2719
+ return cls(id=key_aspect.id)
2273
2720
 
2274
2721
  @property
2275
2722
  def id(self) -> str:
2276
2723
  return self.entity_ids[0]
2277
2724
 
2278
- @property
2279
- def entity_type(self) -> str:
2280
- return self.entity_ids[1]
2281
-
2282
2725
  if TYPE_CHECKING:
2283
- from datahub.metadata.schema_classes import DataHubRetentionKeyClass
2726
+ from datahub.metadata.schema_classes import QueryKeyClass
2284
2727
 
2285
- class DataHubRetentionUrn(_SpecificUrn):
2286
- ENTITY_TYPE: ClassVar[str] = "dataHubRetention"
2287
- _URN_PARTS: ClassVar[int] = 2
2728
+ class QueryUrn(_SpecificUrn):
2729
+ ENTITY_TYPE: ClassVar[str] = "query"
2730
+ _URN_PARTS: ClassVar[int] = 1
2288
2731
 
2289
- def __init__(self, entity_name: str, aspect_name: str, *, _allow_coercion: bool = True) -> None:
2732
+ def __init__(self, id: Union["QueryUrn", str], *, _allow_coercion: bool = True) -> None:
2290
2733
  if _allow_coercion:
2291
2734
  # Field coercion logic (if any is required).
2292
- entity_name = UrnEncoder.encode_string(entity_name)
2293
- aspect_name = UrnEncoder.encode_string(aspect_name)
2735
+ if isinstance(id, str):
2736
+ if id.startswith('urn:li:'):
2737
+ try:
2738
+ id = QueryUrn.from_string(id)
2739
+ except InvalidUrnError:
2740
+ raise InvalidUrnError(f'Expecting a QueryUrn but got {id}')
2741
+ else:
2742
+ id = UrnEncoder.encode_string(id)
2294
2743
 
2295
2744
  # Validation logic.
2296
- if not entity_name:
2297
- raise InvalidUrnError("DataHubRetentionUrn entity_name cannot be empty")
2298
- if UrnEncoder.contains_reserved_char(entity_name):
2299
- raise InvalidUrnError(f'DataHubRetentionUrn entity_name contains reserved characters')
2300
- if not aspect_name:
2301
- raise InvalidUrnError("DataHubRetentionUrn aspect_name cannot be empty")
2302
- if UrnEncoder.contains_reserved_char(aspect_name):
2303
- raise InvalidUrnError(f'DataHubRetentionUrn aspect_name contains reserved characters')
2745
+ if not id:
2746
+ raise InvalidUrnError("QueryUrn id cannot be empty")
2747
+ if isinstance(id, QueryUrn):
2748
+ id = id.id
2749
+ elif isinstance(id, Urn):
2750
+ raise InvalidUrnError(f'Expecting a QueryUrn but got {id}')
2751
+ if UrnEncoder.contains_reserved_char(id):
2752
+ raise InvalidUrnError(f'QueryUrn id contains reserved characters')
2304
2753
 
2305
- super().__init__(self.ENTITY_TYPE, [entity_name, aspect_name])
2754
+ super().__init__(self.ENTITY_TYPE, [id])
2306
2755
 
2307
2756
  @classmethod
2308
- def _parse_ids(cls, entity_ids: List[str]) -> "DataHubRetentionUrn":
2757
+ def _parse_ids(cls, entity_ids: List[str]) -> "QueryUrn":
2309
2758
  if len(entity_ids) != cls._URN_PARTS:
2310
- raise InvalidUrnError(f"DataHubRetentionUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2311
- return cls(entity_name=entity_ids[0], aspect_name=entity_ids[1], _allow_coercion=False)
2759
+ raise InvalidUrnError(f"QueryUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2760
+ return cls(id=entity_ids[0], _allow_coercion=False)
2312
2761
 
2313
2762
  @classmethod
2314
- def underlying_key_aspect_type(cls) -> Type["DataHubRetentionKeyClass"]:
2315
- from datahub.metadata.schema_classes import DataHubRetentionKeyClass
2763
+ def underlying_key_aspect_type(cls) -> Type["QueryKeyClass"]:
2764
+ from datahub.metadata.schema_classes import QueryKeyClass
2316
2765
 
2317
- return DataHubRetentionKeyClass
2766
+ return QueryKeyClass
2318
2767
 
2319
- def to_key_aspect(self) -> "DataHubRetentionKeyClass":
2320
- from datahub.metadata.schema_classes import DataHubRetentionKeyClass
2768
+ def to_key_aspect(self) -> "QueryKeyClass":
2769
+ from datahub.metadata.schema_classes import QueryKeyClass
2321
2770
 
2322
- return DataHubRetentionKeyClass(entityName=self.entity_name, aspectName=self.aspect_name)
2771
+ return QueryKeyClass(id=self.id)
2323
2772
 
2324
2773
  @classmethod
2325
- def from_key_aspect(cls, key_aspect: "DataHubRetentionKeyClass") -> "DataHubRetentionUrn":
2326
- return cls(entity_name=key_aspect.entityName, aspect_name=key_aspect.aspectName)
2774
+ def from_key_aspect(cls, key_aspect: "QueryKeyClass") -> "QueryUrn":
2775
+ return cls(id=key_aspect.id)
2327
2776
 
2328
2777
  @property
2329
- def entity_name(self) -> str:
2778
+ def id(self) -> str:
2330
2779
  return self.entity_ids[0]
2331
2780
 
2332
- @property
2333
- def aspect_name(self) -> str:
2334
- return self.entity_ids[1]
2335
-
2336
2781
  if TYPE_CHECKING:
2337
- from datahub.metadata.schema_classes import TestKeyClass
2782
+ from datahub.metadata.schema_classes import DataHubPersonaKeyClass
2338
2783
 
2339
- class TestUrn(_SpecificUrn):
2340
- ENTITY_TYPE: ClassVar[str] = "test"
2784
+ class DataHubPersonaUrn(_SpecificUrn):
2785
+ ENTITY_TYPE: ClassVar[str] = "dataHubPersona"
2341
2786
  _URN_PARTS: ClassVar[int] = 1
2342
2787
 
2343
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2788
+ def __init__(self, id: Union["DataHubPersonaUrn", str], *, _allow_coercion: bool = True) -> None:
2344
2789
  if _allow_coercion:
2345
2790
  # Field coercion logic (if any is required).
2346
- id = UrnEncoder.encode_string(id)
2791
+ if isinstance(id, str):
2792
+ if id.startswith('urn:li:'):
2793
+ try:
2794
+ id = DataHubPersonaUrn.from_string(id)
2795
+ except InvalidUrnError:
2796
+ raise InvalidUrnError(f'Expecting a DataHubPersonaUrn but got {id}')
2797
+ else:
2798
+ id = UrnEncoder.encode_string(id)
2347
2799
 
2348
2800
  # Validation logic.
2349
2801
  if not id:
2350
- raise InvalidUrnError("TestUrn id cannot be empty")
2802
+ raise InvalidUrnError("DataHubPersonaUrn id cannot be empty")
2803
+ if isinstance(id, DataHubPersonaUrn):
2804
+ id = id.id
2805
+ elif isinstance(id, Urn):
2806
+ raise InvalidUrnError(f'Expecting a DataHubPersonaUrn but got {id}')
2351
2807
  if UrnEncoder.contains_reserved_char(id):
2352
- raise InvalidUrnError(f'TestUrn id contains reserved characters')
2808
+ raise InvalidUrnError(f'DataHubPersonaUrn id contains reserved characters')
2353
2809
 
2354
2810
  super().__init__(self.ENTITY_TYPE, [id])
2355
2811
 
2356
2812
  @classmethod
2357
- def _parse_ids(cls, entity_ids: List[str]) -> "TestUrn":
2813
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubPersonaUrn":
2358
2814
  if len(entity_ids) != cls._URN_PARTS:
2359
- raise InvalidUrnError(f"TestUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2815
+ raise InvalidUrnError(f"DataHubPersonaUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2360
2816
  return cls(id=entity_ids[0], _allow_coercion=False)
2361
2817
 
2362
2818
  @classmethod
2363
- def underlying_key_aspect_type(cls) -> Type["TestKeyClass"]:
2364
- from datahub.metadata.schema_classes import TestKeyClass
2819
+ def underlying_key_aspect_type(cls) -> Type["DataHubPersonaKeyClass"]:
2820
+ from datahub.metadata.schema_classes import DataHubPersonaKeyClass
2365
2821
 
2366
- return TestKeyClass
2822
+ return DataHubPersonaKeyClass
2367
2823
 
2368
- def to_key_aspect(self) -> "TestKeyClass":
2369
- from datahub.metadata.schema_classes import TestKeyClass
2824
+ def to_key_aspect(self) -> "DataHubPersonaKeyClass":
2825
+ from datahub.metadata.schema_classes import DataHubPersonaKeyClass
2370
2826
 
2371
- return TestKeyClass(id=self.id)
2827
+ return DataHubPersonaKeyClass(id=self.id)
2372
2828
 
2373
2829
  @classmethod
2374
- def from_key_aspect(cls, key_aspect: "TestKeyClass") -> "TestUrn":
2830
+ def from_key_aspect(cls, key_aspect: "DataHubPersonaKeyClass") -> "DataHubPersonaUrn":
2375
2831
  return cls(id=key_aspect.id)
2376
2832
 
2377
2833
  @property
@@ -2433,383 +2889,354 @@ class ChartUrn(_SpecificUrn):
2433
2889
  return self.entity_ids[1]
2434
2890
 
2435
2891
  if TYPE_CHECKING:
2436
- from datahub.metadata.schema_classes import MLModelGroupKeyClass
2892
+ from datahub.metadata.schema_classes import VersionSetKeyClass
2437
2893
 
2438
- class MlModelGroupUrn(_SpecificUrn):
2439
- ENTITY_TYPE: ClassVar[str] = "mlModelGroup"
2440
- _URN_PARTS: ClassVar[int] = 3
2894
+ class VersionSetUrn(_SpecificUrn):
2895
+ ENTITY_TYPE: ClassVar[str] = "versionSet"
2896
+ _URN_PARTS: ClassVar[int] = 2
2441
2897
 
2442
- def __init__(self, platform: Union["DataPlatformUrn", str], name: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
2898
+ def __init__(self, id: str, entity_type: str, *, _allow_coercion: bool = True) -> None:
2443
2899
  if _allow_coercion:
2444
2900
  # Field coercion logic (if any is required).
2445
- platform = platform.urn() if isinstance(platform, DataPlatformUrn) else DataPlatformUrn(platform).urn()
2446
- name = UrnEncoder.encode_string(name)
2447
- env = env.upper()
2448
- env = UrnEncoder.encode_string(env)
2901
+ id = UrnEncoder.encode_string(id)
2902
+ entity_type = UrnEncoder.encode_string(entity_type)
2449
2903
 
2450
2904
  # Validation logic.
2451
- if not platform:
2452
- raise InvalidUrnError("MlModelGroupUrn platform cannot be empty")
2453
- platform = str(platform)
2454
- assert DataPlatformUrn.from_string(platform)
2455
- if not name:
2456
- raise InvalidUrnError("MlModelGroupUrn name cannot be empty")
2457
- if UrnEncoder.contains_reserved_char(name):
2458
- raise InvalidUrnError(f'MlModelGroupUrn name contains reserved characters')
2459
- if not env:
2460
- raise InvalidUrnError("MlModelGroupUrn env cannot be empty")
2461
- if UrnEncoder.contains_reserved_char(env):
2462
- raise InvalidUrnError(f'MlModelGroupUrn env contains reserved characters')
2905
+ if not id:
2906
+ raise InvalidUrnError("VersionSetUrn id cannot be empty")
2907
+ if UrnEncoder.contains_reserved_char(id):
2908
+ raise InvalidUrnError(f'VersionSetUrn id contains reserved characters')
2909
+ if not entity_type:
2910
+ raise InvalidUrnError("VersionSetUrn entity_type cannot be empty")
2911
+ if UrnEncoder.contains_reserved_char(entity_type):
2912
+ raise InvalidUrnError(f'VersionSetUrn entity_type contains reserved characters')
2463
2913
 
2464
- super().__init__(self.ENTITY_TYPE, [platform, name, env])
2914
+ super().__init__(self.ENTITY_TYPE, [id, entity_type])
2465
2915
 
2466
2916
  @classmethod
2467
- def _parse_ids(cls, entity_ids: List[str]) -> "MlModelGroupUrn":
2917
+ def _parse_ids(cls, entity_ids: List[str]) -> "VersionSetUrn":
2468
2918
  if len(entity_ids) != cls._URN_PARTS:
2469
- raise InvalidUrnError(f"MlModelGroupUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2470
- return cls(platform=entity_ids[0], name=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
2919
+ raise InvalidUrnError(f"VersionSetUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2920
+ return cls(id=entity_ids[0], entity_type=entity_ids[1], _allow_coercion=False)
2471
2921
 
2472
2922
  @classmethod
2473
- def underlying_key_aspect_type(cls) -> Type["MLModelGroupKeyClass"]:
2474
- from datahub.metadata.schema_classes import MLModelGroupKeyClass
2923
+ def underlying_key_aspect_type(cls) -> Type["VersionSetKeyClass"]:
2924
+ from datahub.metadata.schema_classes import VersionSetKeyClass
2475
2925
 
2476
- return MLModelGroupKeyClass
2926
+ return VersionSetKeyClass
2477
2927
 
2478
- def to_key_aspect(self) -> "MLModelGroupKeyClass":
2479
- from datahub.metadata.schema_classes import MLModelGroupKeyClass
2928
+ def to_key_aspect(self) -> "VersionSetKeyClass":
2929
+ from datahub.metadata.schema_classes import VersionSetKeyClass
2480
2930
 
2481
- return MLModelGroupKeyClass(platform=self.platform, name=self.name, origin=self.env)
2931
+ return VersionSetKeyClass(id=self.id, entityType=self.entity_type)
2482
2932
 
2483
2933
  @classmethod
2484
- def from_key_aspect(cls, key_aspect: "MLModelGroupKeyClass") -> "MlModelGroupUrn":
2485
- return cls(platform=key_aspect.platform, name=key_aspect.name, env=key_aspect.origin)
2934
+ def from_key_aspect(cls, key_aspect: "VersionSetKeyClass") -> "VersionSetUrn":
2935
+ return cls(id=key_aspect.id, entity_type=key_aspect.entityType)
2486
2936
 
2487
2937
  @property
2488
- def platform(self) -> str:
2938
+ def id(self) -> str:
2489
2939
  return self.entity_ids[0]
2490
2940
 
2491
2941
  @property
2492
- def name(self) -> str:
2942
+ def entity_type(self) -> str:
2493
2943
  return self.entity_ids[1]
2494
2944
 
2495
- @property
2496
- def env(self) -> str:
2497
- return self.entity_ids[2]
2498
-
2499
- if TYPE_CHECKING:
2500
- from datahub.metadata.schema_classes import TelemetryKeyClass
2501
-
2502
- class TelemetryUrn(_SpecificUrn):
2503
- ENTITY_TYPE: ClassVar[str] = "telemetry"
2504
- _URN_PARTS: ClassVar[int] = 1
2505
-
2506
- def __init__(self, name: str, *, _allow_coercion: bool = True) -> None:
2507
- if _allow_coercion:
2508
- # Field coercion logic (if any is required).
2509
- name = UrnEncoder.encode_string(name)
2510
-
2511
- # Validation logic.
2512
- if not name:
2513
- raise InvalidUrnError("TelemetryUrn name cannot be empty")
2514
- if UrnEncoder.contains_reserved_char(name):
2515
- raise InvalidUrnError(f'TelemetryUrn name contains reserved characters')
2516
-
2517
- super().__init__(self.ENTITY_TYPE, [name])
2518
-
2519
- @classmethod
2520
- def _parse_ids(cls, entity_ids: List[str]) -> "TelemetryUrn":
2521
- if len(entity_ids) != cls._URN_PARTS:
2522
- raise InvalidUrnError(f"TelemetryUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2523
- return cls(name=entity_ids[0], _allow_coercion=False)
2524
-
2525
- @classmethod
2526
- def underlying_key_aspect_type(cls) -> Type["TelemetryKeyClass"]:
2527
- from datahub.metadata.schema_classes import TelemetryKeyClass
2528
-
2529
- return TelemetryKeyClass
2530
-
2531
- def to_key_aspect(self) -> "TelemetryKeyClass":
2532
- from datahub.metadata.schema_classes import TelemetryKeyClass
2533
-
2534
- return TelemetryKeyClass(name=self.name)
2535
-
2536
- @classmethod
2537
- def from_key_aspect(cls, key_aspect: "TelemetryKeyClass") -> "TelemetryUrn":
2538
- return cls(name=key_aspect.name)
2539
-
2540
- @property
2541
- def name(self) -> str:
2542
- return self.entity_ids[0]
2543
-
2544
2945
  if TYPE_CHECKING:
2545
- from datahub.metadata.schema_classes import DomainKeyClass
2946
+ from datahub.metadata.schema_classes import DataHubAccessTokenKeyClass
2546
2947
 
2547
- class DomainUrn(_SpecificUrn):
2548
- ENTITY_TYPE: ClassVar[str] = "domain"
2948
+ class DataHubAccessTokenUrn(_SpecificUrn):
2949
+ ENTITY_TYPE: ClassVar[str] = "dataHubAccessToken"
2549
2950
  _URN_PARTS: ClassVar[int] = 1
2550
2951
 
2551
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2952
+ def __init__(self, id: Union["DataHubAccessTokenUrn", str], *, _allow_coercion: bool = True) -> None:
2552
2953
  if _allow_coercion:
2553
2954
  # Field coercion logic (if any is required).
2554
- id = UrnEncoder.encode_string(id)
2955
+ if isinstance(id, str):
2956
+ if id.startswith('urn:li:'):
2957
+ try:
2958
+ id = DataHubAccessTokenUrn.from_string(id)
2959
+ except InvalidUrnError:
2960
+ raise InvalidUrnError(f'Expecting a DataHubAccessTokenUrn but got {id}')
2961
+ else:
2962
+ id = UrnEncoder.encode_string(id)
2555
2963
 
2556
2964
  # Validation logic.
2557
2965
  if not id:
2558
- raise InvalidUrnError("DomainUrn id cannot be empty")
2966
+ raise InvalidUrnError("DataHubAccessTokenUrn id cannot be empty")
2967
+ if isinstance(id, DataHubAccessTokenUrn):
2968
+ id = id.id
2969
+ elif isinstance(id, Urn):
2970
+ raise InvalidUrnError(f'Expecting a DataHubAccessTokenUrn but got {id}')
2559
2971
  if UrnEncoder.contains_reserved_char(id):
2560
- raise InvalidUrnError(f'DomainUrn id contains reserved characters')
2972
+ raise InvalidUrnError(f'DataHubAccessTokenUrn id contains reserved characters')
2561
2973
 
2562
2974
  super().__init__(self.ENTITY_TYPE, [id])
2563
2975
 
2564
2976
  @classmethod
2565
- def _parse_ids(cls, entity_ids: List[str]) -> "DomainUrn":
2977
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubAccessTokenUrn":
2566
2978
  if len(entity_ids) != cls._URN_PARTS:
2567
- raise InvalidUrnError(f"DomainUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2979
+ raise InvalidUrnError(f"DataHubAccessTokenUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2568
2980
  return cls(id=entity_ids[0], _allow_coercion=False)
2569
2981
 
2570
2982
  @classmethod
2571
- def underlying_key_aspect_type(cls) -> Type["DomainKeyClass"]:
2572
- from datahub.metadata.schema_classes import DomainKeyClass
2983
+ def underlying_key_aspect_type(cls) -> Type["DataHubAccessTokenKeyClass"]:
2984
+ from datahub.metadata.schema_classes import DataHubAccessTokenKeyClass
2573
2985
 
2574
- return DomainKeyClass
2986
+ return DataHubAccessTokenKeyClass
2575
2987
 
2576
- def to_key_aspect(self) -> "DomainKeyClass":
2577
- from datahub.metadata.schema_classes import DomainKeyClass
2988
+ def to_key_aspect(self) -> "DataHubAccessTokenKeyClass":
2989
+ from datahub.metadata.schema_classes import DataHubAccessTokenKeyClass
2578
2990
 
2579
- return DomainKeyClass(id=self.id)
2991
+ return DataHubAccessTokenKeyClass(id=self.id)
2580
2992
 
2581
2993
  @classmethod
2582
- def from_key_aspect(cls, key_aspect: "DomainKeyClass") -> "DomainUrn":
2994
+ def from_key_aspect(cls, key_aspect: "DataHubAccessTokenKeyClass") -> "DataHubAccessTokenUrn":
2583
2995
  return cls(id=key_aspect.id)
2584
2996
 
2585
- @classmethod
2586
- @deprecated(reason="Use the constructor instead")
2587
- def create_from_id(cls, id: str) -> "DomainUrn":
2588
- return cls(id)
2589
-
2590
2997
  @property
2591
2998
  def id(self) -> str:
2592
2999
  return self.entity_ids[0]
2593
3000
 
2594
3001
  if TYPE_CHECKING:
2595
- from datahub.metadata.schema_classes import IncidentKeyClass
3002
+ from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
2596
3003
 
2597
- class IncidentUrn(_SpecificUrn):
2598
- ENTITY_TYPE: ClassVar[str] = "incident"
3004
+ class DataProcessInstanceUrn(_SpecificUrn):
3005
+ ENTITY_TYPE: ClassVar[str] = "dataProcessInstance"
2599
3006
  _URN_PARTS: ClassVar[int] = 1
2600
3007
 
2601
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
3008
+ def __init__(self, id: Union["DataProcessInstanceUrn", str], *, _allow_coercion: bool = True) -> None:
2602
3009
  if _allow_coercion:
2603
3010
  # Field coercion logic (if any is required).
2604
- id = UrnEncoder.encode_string(id)
3011
+ if isinstance(id, str):
3012
+ if id.startswith('urn:li:'):
3013
+ try:
3014
+ id = DataProcessInstanceUrn.from_string(id)
3015
+ except InvalidUrnError:
3016
+ raise InvalidUrnError(f'Expecting a DataProcessInstanceUrn but got {id}')
3017
+ else:
3018
+ id = UrnEncoder.encode_string(id)
2605
3019
 
2606
3020
  # Validation logic.
2607
3021
  if not id:
2608
- raise InvalidUrnError("IncidentUrn id cannot be empty")
3022
+ raise InvalidUrnError("DataProcessInstanceUrn id cannot be empty")
3023
+ if isinstance(id, DataProcessInstanceUrn):
3024
+ id = id.id
3025
+ elif isinstance(id, Urn):
3026
+ raise InvalidUrnError(f'Expecting a DataProcessInstanceUrn but got {id}')
2609
3027
  if UrnEncoder.contains_reserved_char(id):
2610
- raise InvalidUrnError(f'IncidentUrn id contains reserved characters')
3028
+ raise InvalidUrnError(f'DataProcessInstanceUrn id contains reserved characters')
2611
3029
 
2612
3030
  super().__init__(self.ENTITY_TYPE, [id])
2613
3031
 
2614
3032
  @classmethod
2615
- def _parse_ids(cls, entity_ids: List[str]) -> "IncidentUrn":
3033
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataProcessInstanceUrn":
2616
3034
  if len(entity_ids) != cls._URN_PARTS:
2617
- raise InvalidUrnError(f"IncidentUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
3035
+ raise InvalidUrnError(f"DataProcessInstanceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2618
3036
  return cls(id=entity_ids[0], _allow_coercion=False)
2619
3037
 
2620
3038
  @classmethod
2621
- def underlying_key_aspect_type(cls) -> Type["IncidentKeyClass"]:
2622
- from datahub.metadata.schema_classes import IncidentKeyClass
3039
+ def underlying_key_aspect_type(cls) -> Type["DataProcessInstanceKeyClass"]:
3040
+ from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
2623
3041
 
2624
- return IncidentKeyClass
3042
+ return DataProcessInstanceKeyClass
2625
3043
 
2626
- def to_key_aspect(self) -> "IncidentKeyClass":
2627
- from datahub.metadata.schema_classes import IncidentKeyClass
3044
+ def to_key_aspect(self) -> "DataProcessInstanceKeyClass":
3045
+ from datahub.metadata.schema_classes import DataProcessInstanceKeyClass
2628
3046
 
2629
- return IncidentKeyClass(id=self.id)
3047
+ return DataProcessInstanceKeyClass(id=self.id)
2630
3048
 
2631
3049
  @classmethod
2632
- def from_key_aspect(cls, key_aspect: "IncidentKeyClass") -> "IncidentUrn":
3050
+ def from_key_aspect(cls, key_aspect: "DataProcessInstanceKeyClass") -> "DataProcessInstanceUrn":
2633
3051
  return cls(id=key_aspect.id)
2634
3052
 
3053
+ @classmethod
3054
+ @deprecated(reason="Use the constructor instead")
3055
+ def create_from_id(cls, id: str) -> "DataProcessInstanceUrn":
3056
+ return cls(id)
3057
+
3058
+ @deprecated(reason="Use .id instead")
3059
+ def get_dataprocessinstance_id(self) -> str:
3060
+ return self.id
3061
+
2635
3062
  @property
2636
3063
  def id(self) -> str:
2637
3064
  return self.entity_ids[0]
2638
3065
 
2639
3066
  if TYPE_CHECKING:
2640
- from datahub.metadata.schema_classes import GlossaryTermKeyClass
3067
+ from datahub.metadata.schema_classes import CorpGroupKeyClass
2641
3068
 
2642
- class GlossaryTermUrn(_SpecificUrn):
2643
- ENTITY_TYPE: ClassVar[str] = "glossaryTerm"
3069
+ class CorpGroupUrn(_SpecificUrn):
3070
+ ENTITY_TYPE: ClassVar[str] = "corpGroup"
2644
3071
  _URN_PARTS: ClassVar[int] = 1
2645
3072
 
2646
- def __init__(self, name: str, *, _allow_coercion: bool = True) -> None:
3073
+ def __init__(self, name: Union["CorpGroupUrn", str], *, _allow_coercion: bool = True) -> None:
2647
3074
  if _allow_coercion:
2648
3075
  # Field coercion logic (if any is required).
2649
- name = UrnEncoder.encode_string(name)
3076
+ if isinstance(name, str):
3077
+ if name.startswith('urn:li:'):
3078
+ try:
3079
+ name = CorpGroupUrn.from_string(name)
3080
+ except InvalidUrnError:
3081
+ raise InvalidUrnError(f'Expecting a CorpGroupUrn but got {name}')
3082
+ else:
3083
+ name = UrnEncoder.encode_string(name)
2650
3084
 
2651
3085
  # Validation logic.
2652
3086
  if not name:
2653
- raise InvalidUrnError("GlossaryTermUrn name cannot be empty")
3087
+ raise InvalidUrnError("CorpGroupUrn name cannot be empty")
3088
+ if isinstance(name, CorpGroupUrn):
3089
+ name = name.name
3090
+ elif isinstance(name, Urn):
3091
+ raise InvalidUrnError(f'Expecting a CorpGroupUrn but got {name}')
2654
3092
  if UrnEncoder.contains_reserved_char(name):
2655
- raise InvalidUrnError(f'GlossaryTermUrn name contains reserved characters')
3093
+ raise InvalidUrnError(f'CorpGroupUrn name contains reserved characters')
2656
3094
 
2657
3095
  super().__init__(self.ENTITY_TYPE, [name])
2658
3096
 
2659
3097
  @classmethod
2660
- def _parse_ids(cls, entity_ids: List[str]) -> "GlossaryTermUrn":
3098
+ def _parse_ids(cls, entity_ids: List[str]) -> "CorpGroupUrn":
2661
3099
  if len(entity_ids) != cls._URN_PARTS:
2662
- raise InvalidUrnError(f"GlossaryTermUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
3100
+ raise InvalidUrnError(f"CorpGroupUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2663
3101
  return cls(name=entity_ids[0], _allow_coercion=False)
2664
3102
 
2665
3103
  @classmethod
2666
- def underlying_key_aspect_type(cls) -> Type["GlossaryTermKeyClass"]:
2667
- from datahub.metadata.schema_classes import GlossaryTermKeyClass
3104
+ def underlying_key_aspect_type(cls) -> Type["CorpGroupKeyClass"]:
3105
+ from datahub.metadata.schema_classes import CorpGroupKeyClass
2668
3106
 
2669
- return GlossaryTermKeyClass
3107
+ return CorpGroupKeyClass
2670
3108
 
2671
- def to_key_aspect(self) -> "GlossaryTermKeyClass":
2672
- from datahub.metadata.schema_classes import GlossaryTermKeyClass
3109
+ def to_key_aspect(self) -> "CorpGroupKeyClass":
3110
+ from datahub.metadata.schema_classes import CorpGroupKeyClass
2673
3111
 
2674
- return GlossaryTermKeyClass(name=self.name)
3112
+ return CorpGroupKeyClass(name=self.name)
2675
3113
 
2676
3114
  @classmethod
2677
- def from_key_aspect(cls, key_aspect: "GlossaryTermKeyClass") -> "GlossaryTermUrn":
3115
+ def from_key_aspect(cls, key_aspect: "CorpGroupKeyClass") -> "CorpGroupUrn":
2678
3116
  return cls(name=key_aspect.name)
2679
3117
 
3118
+ @classmethod
3119
+ @deprecated(reason="Use the constructor instead")
3120
+ def create_from_id(cls, id: str) -> "CorpGroupUrn":
3121
+ return cls(id)
3122
+
2680
3123
  @property
2681
3124
  def name(self) -> str:
2682
3125
  return self.entity_ids[0]
2683
3126
 
2684
3127
  if TYPE_CHECKING:
2685
- from datahub.metadata.schema_classes import InviteTokenKeyClass
3128
+ from datahub.metadata.schema_classes import DataProcessKeyClass
2686
3129
 
2687
- class InviteTokenUrn(_SpecificUrn):
2688
- ENTITY_TYPE: ClassVar[str] = "inviteToken"
2689
- _URN_PARTS: ClassVar[int] = 1
3130
+ class DataProcessUrn(_SpecificUrn):
3131
+ ENTITY_TYPE: ClassVar[str] = "dataProcess"
3132
+ _URN_PARTS: ClassVar[int] = 3
2690
3133
 
2691
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
3134
+ def __init__(self, name: str, orchestrator: str, env: str = "PROD", *, _allow_coercion: bool = True) -> None:
2692
3135
  if _allow_coercion:
2693
3136
  # Field coercion logic (if any is required).
2694
- id = UrnEncoder.encode_string(id)
3137
+ name = UrnEncoder.encode_string(name)
3138
+ orchestrator = UrnEncoder.encode_string(orchestrator)
3139
+ env = env.upper()
2695
3140
 
2696
3141
  # Validation logic.
2697
- if not id:
2698
- raise InvalidUrnError("InviteTokenUrn id cannot be empty")
2699
- if UrnEncoder.contains_reserved_char(id):
2700
- raise InvalidUrnError(f'InviteTokenUrn id contains reserved characters')
3142
+ if not name:
3143
+ raise InvalidUrnError("DataProcessUrn name cannot be empty")
3144
+ if UrnEncoder.contains_reserved_char(name):
3145
+ raise InvalidUrnError(f'DataProcessUrn name contains reserved characters')
3146
+ if not orchestrator:
3147
+ raise InvalidUrnError("DataProcessUrn orchestrator cannot be empty")
3148
+ if UrnEncoder.contains_reserved_char(orchestrator):
3149
+ raise InvalidUrnError(f'DataProcessUrn orchestrator contains reserved characters')
3150
+ if not env:
3151
+ raise InvalidUrnError("DataProcessUrn env cannot be empty")
3152
+ if UrnEncoder.contains_reserved_char(env):
3153
+ raise InvalidUrnError(f'DataProcessUrn env contains reserved characters')
2701
3154
 
2702
- super().__init__(self.ENTITY_TYPE, [id])
3155
+ super().__init__(self.ENTITY_TYPE, [name, orchestrator, env])
2703
3156
 
2704
3157
  @classmethod
2705
- def _parse_ids(cls, entity_ids: List[str]) -> "InviteTokenUrn":
3158
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataProcessUrn":
2706
3159
  if len(entity_ids) != cls._URN_PARTS:
2707
- raise InvalidUrnError(f"InviteTokenUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2708
- return cls(id=entity_ids[0], _allow_coercion=False)
3160
+ raise InvalidUrnError(f"DataProcessUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
3161
+ return cls(name=entity_ids[0], orchestrator=entity_ids[1], env=entity_ids[2], _allow_coercion=False)
2709
3162
 
2710
3163
  @classmethod
2711
- def underlying_key_aspect_type(cls) -> Type["InviteTokenKeyClass"]:
2712
- from datahub.metadata.schema_classes import InviteTokenKeyClass
3164
+ def underlying_key_aspect_type(cls) -> Type["DataProcessKeyClass"]:
3165
+ from datahub.metadata.schema_classes import DataProcessKeyClass
2713
3166
 
2714
- return InviteTokenKeyClass
3167
+ return DataProcessKeyClass
2715
3168
 
2716
- def to_key_aspect(self) -> "InviteTokenKeyClass":
2717
- from datahub.metadata.schema_classes import InviteTokenKeyClass
3169
+ def to_key_aspect(self) -> "DataProcessKeyClass":
3170
+ from datahub.metadata.schema_classes import DataProcessKeyClass
2718
3171
 
2719
- return InviteTokenKeyClass(id=self.id)
3172
+ return DataProcessKeyClass(name=self.name, orchestrator=self.orchestrator, origin=self.env)
2720
3173
 
2721
3174
  @classmethod
2722
- def from_key_aspect(cls, key_aspect: "InviteTokenKeyClass") -> "InviteTokenUrn":
2723
- return cls(id=key_aspect.id)
3175
+ def from_key_aspect(cls, key_aspect: "DataProcessKeyClass") -> "DataProcessUrn":
3176
+ return cls(name=key_aspect.name, orchestrator=key_aspect.orchestrator, env=key_aspect.origin)
2724
3177
 
2725
3178
  @property
2726
- def id(self) -> str:
3179
+ def name(self) -> str:
2727
3180
  return self.entity_ids[0]
2728
3181
 
2729
- if TYPE_CHECKING:
2730
- from datahub.metadata.schema_classes import PlatformResourceKeyClass
2731
-
2732
- class PlatformResourceUrn(_SpecificUrn):
2733
- ENTITY_TYPE: ClassVar[str] = "platformResource"
2734
- _URN_PARTS: ClassVar[int] = 1
2735
-
2736
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
2737
- if _allow_coercion:
2738
- # Field coercion logic (if any is required).
2739
- id = UrnEncoder.encode_string(id)
2740
-
2741
- # Validation logic.
2742
- if not id:
2743
- raise InvalidUrnError("PlatformResourceUrn id cannot be empty")
2744
- if UrnEncoder.contains_reserved_char(id):
2745
- raise InvalidUrnError(f'PlatformResourceUrn id contains reserved characters')
2746
-
2747
- super().__init__(self.ENTITY_TYPE, [id])
2748
-
2749
- @classmethod
2750
- def _parse_ids(cls, entity_ids: List[str]) -> "PlatformResourceUrn":
2751
- if len(entity_ids) != cls._URN_PARTS:
2752
- raise InvalidUrnError(f"PlatformResourceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2753
- return cls(id=entity_ids[0], _allow_coercion=False)
2754
-
2755
- @classmethod
2756
- def underlying_key_aspect_type(cls) -> Type["PlatformResourceKeyClass"]:
2757
- from datahub.metadata.schema_classes import PlatformResourceKeyClass
2758
-
2759
- return PlatformResourceKeyClass
2760
-
2761
- def to_key_aspect(self) -> "PlatformResourceKeyClass":
2762
- from datahub.metadata.schema_classes import PlatformResourceKeyClass
2763
-
2764
- return PlatformResourceKeyClass(id=self.id)
2765
-
2766
- @classmethod
2767
- def from_key_aspect(cls, key_aspect: "PlatformResourceKeyClass") -> "PlatformResourceUrn":
2768
- return cls(id=key_aspect.id)
3182
+ @property
3183
+ def orchestrator(self) -> str:
3184
+ return self.entity_ids[1]
2769
3185
 
2770
3186
  @property
2771
- def id(self) -> str:
2772
- return self.entity_ids[0]
3187
+ def env(self) -> str:
3188
+ return self.entity_ids[2]
2773
3189
 
2774
3190
  if TYPE_CHECKING:
2775
- from datahub.metadata.schema_classes import EntityTypeKeyClass
3191
+ from datahub.metadata.schema_classes import DataHubIngestionSourceKeyClass
2776
3192
 
2777
- class EntityTypeUrn(_SpecificUrn):
2778
- ENTITY_TYPE: ClassVar[str] = "entityType"
3193
+ class DataHubIngestionSourceUrn(_SpecificUrn):
3194
+ ENTITY_TYPE: ClassVar[str] = "dataHubIngestionSource"
2779
3195
  _URN_PARTS: ClassVar[int] = 1
2780
3196
 
2781
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
3197
+ def __init__(self, id: Union["DataHubIngestionSourceUrn", str], *, _allow_coercion: bool = True) -> None:
2782
3198
  if _allow_coercion:
2783
3199
  # Field coercion logic (if any is required).
2784
- id = UrnEncoder.encode_string(id)
3200
+ if isinstance(id, str):
3201
+ if id.startswith('urn:li:'):
3202
+ try:
3203
+ id = DataHubIngestionSourceUrn.from_string(id)
3204
+ except InvalidUrnError:
3205
+ raise InvalidUrnError(f'Expecting a DataHubIngestionSourceUrn but got {id}')
3206
+ else:
3207
+ id = UrnEncoder.encode_string(id)
2785
3208
 
2786
3209
  # Validation logic.
2787
3210
  if not id:
2788
- raise InvalidUrnError("EntityTypeUrn id cannot be empty")
3211
+ raise InvalidUrnError("DataHubIngestionSourceUrn id cannot be empty")
3212
+ if isinstance(id, DataHubIngestionSourceUrn):
3213
+ id = id.id
3214
+ elif isinstance(id, Urn):
3215
+ raise InvalidUrnError(f'Expecting a DataHubIngestionSourceUrn but got {id}')
2789
3216
  if UrnEncoder.contains_reserved_char(id):
2790
- raise InvalidUrnError(f'EntityTypeUrn id contains reserved characters')
3217
+ raise InvalidUrnError(f'DataHubIngestionSourceUrn id contains reserved characters')
2791
3218
 
2792
3219
  super().__init__(self.ENTITY_TYPE, [id])
2793
3220
 
2794
3221
  @classmethod
2795
- def _parse_ids(cls, entity_ids: List[str]) -> "EntityTypeUrn":
3222
+ def _parse_ids(cls, entity_ids: List[str]) -> "DataHubIngestionSourceUrn":
2796
3223
  if len(entity_ids) != cls._URN_PARTS:
2797
- raise InvalidUrnError(f"EntityTypeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
3224
+ raise InvalidUrnError(f"DataHubIngestionSourceUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2798
3225
  return cls(id=entity_ids[0], _allow_coercion=False)
2799
3226
 
2800
3227
  @classmethod
2801
- def underlying_key_aspect_type(cls) -> Type["EntityTypeKeyClass"]:
2802
- from datahub.metadata.schema_classes import EntityTypeKeyClass
3228
+ def underlying_key_aspect_type(cls) -> Type["DataHubIngestionSourceKeyClass"]:
3229
+ from datahub.metadata.schema_classes import DataHubIngestionSourceKeyClass
2803
3230
 
2804
- return EntityTypeKeyClass
3231
+ return DataHubIngestionSourceKeyClass
2805
3232
 
2806
- def to_key_aspect(self) -> "EntityTypeKeyClass":
2807
- from datahub.metadata.schema_classes import EntityTypeKeyClass
3233
+ def to_key_aspect(self) -> "DataHubIngestionSourceKeyClass":
3234
+ from datahub.metadata.schema_classes import DataHubIngestionSourceKeyClass
2808
3235
 
2809
- return EntityTypeKeyClass(id=self.id)
3236
+ return DataHubIngestionSourceKeyClass(id=self.id)
2810
3237
 
2811
3238
  @classmethod
2812
- def from_key_aspect(cls, key_aspect: "EntityTypeKeyClass") -> "EntityTypeUrn":
3239
+ def from_key_aspect(cls, key_aspect: "DataHubIngestionSourceKeyClass") -> "DataHubIngestionSourceUrn":
2813
3240
  return cls(id=key_aspect.id)
2814
3241
 
2815
3242
  @property
@@ -2817,44 +3244,55 @@ class EntityTypeUrn(_SpecificUrn):
2817
3244
  return self.entity_ids[0]
2818
3245
 
2819
3246
  if TYPE_CHECKING:
2820
- from datahub.metadata.schema_classes import DataProductKeyClass
3247
+ from datahub.metadata.schema_classes import GlobalSettingsKeyClass
2821
3248
 
2822
- class DataProductUrn(_SpecificUrn):
2823
- ENTITY_TYPE: ClassVar[str] = "dataProduct"
3249
+ class GlobalSettingsUrn(_SpecificUrn):
3250
+ ENTITY_TYPE: ClassVar[str] = "globalSettings"
2824
3251
  _URN_PARTS: ClassVar[int] = 1
2825
3252
 
2826
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
3253
+ def __init__(self, id: Union["GlobalSettingsUrn", str], *, _allow_coercion: bool = True) -> None:
2827
3254
  if _allow_coercion:
2828
3255
  # Field coercion logic (if any is required).
2829
- id = UrnEncoder.encode_string(id)
3256
+ if isinstance(id, str):
3257
+ if id.startswith('urn:li:'):
3258
+ try:
3259
+ id = GlobalSettingsUrn.from_string(id)
3260
+ except InvalidUrnError:
3261
+ raise InvalidUrnError(f'Expecting a GlobalSettingsUrn but got {id}')
3262
+ else:
3263
+ id = UrnEncoder.encode_string(id)
2830
3264
 
2831
3265
  # Validation logic.
2832
3266
  if not id:
2833
- raise InvalidUrnError("DataProductUrn id cannot be empty")
3267
+ raise InvalidUrnError("GlobalSettingsUrn id cannot be empty")
3268
+ if isinstance(id, GlobalSettingsUrn):
3269
+ id = id.id
3270
+ elif isinstance(id, Urn):
3271
+ raise InvalidUrnError(f'Expecting a GlobalSettingsUrn but got {id}')
2834
3272
  if UrnEncoder.contains_reserved_char(id):
2835
- raise InvalidUrnError(f'DataProductUrn id contains reserved characters')
3273
+ raise InvalidUrnError(f'GlobalSettingsUrn id contains reserved characters')
2836
3274
 
2837
3275
  super().__init__(self.ENTITY_TYPE, [id])
2838
3276
 
2839
3277
  @classmethod
2840
- def _parse_ids(cls, entity_ids: List[str]) -> "DataProductUrn":
3278
+ def _parse_ids(cls, entity_ids: List[str]) -> "GlobalSettingsUrn":
2841
3279
  if len(entity_ids) != cls._URN_PARTS:
2842
- raise InvalidUrnError(f"DataProductUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
3280
+ raise InvalidUrnError(f"GlobalSettingsUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2843
3281
  return cls(id=entity_ids[0], _allow_coercion=False)
2844
3282
 
2845
3283
  @classmethod
2846
- def underlying_key_aspect_type(cls) -> Type["DataProductKeyClass"]:
2847
- from datahub.metadata.schema_classes import DataProductKeyClass
3284
+ def underlying_key_aspect_type(cls) -> Type["GlobalSettingsKeyClass"]:
3285
+ from datahub.metadata.schema_classes import GlobalSettingsKeyClass
2848
3286
 
2849
- return DataProductKeyClass
3287
+ return GlobalSettingsKeyClass
2850
3288
 
2851
- def to_key_aspect(self) -> "DataProductKeyClass":
2852
- from datahub.metadata.schema_classes import DataProductKeyClass
3289
+ def to_key_aspect(self) -> "GlobalSettingsKeyClass":
3290
+ from datahub.metadata.schema_classes import GlobalSettingsKeyClass
2853
3291
 
2854
- return DataProductKeyClass(id=self.id)
3292
+ return GlobalSettingsKeyClass(id=self.id)
2855
3293
 
2856
3294
  @classmethod
2857
- def from_key_aspect(cls, key_aspect: "DataProductKeyClass") -> "DataProductUrn":
3295
+ def from_key_aspect(cls, key_aspect: "GlobalSettingsKeyClass") -> "GlobalSettingsUrn":
2858
3296
  return cls(id=key_aspect.id)
2859
3297
 
2860
3298
  @property
@@ -2862,44 +3300,55 @@ class DataProductUrn(_SpecificUrn):
2862
3300
  return self.entity_ids[0]
2863
3301
 
2864
3302
  if TYPE_CHECKING:
2865
- from datahub.metadata.schema_classes import BusinessAttributeKeyClass
3303
+ from datahub.metadata.schema_classes import EntityTypeKeyClass
2866
3304
 
2867
- class BusinessAttributeUrn(_SpecificUrn):
2868
- ENTITY_TYPE: ClassVar[str] = "businessAttribute"
3305
+ class EntityTypeUrn(_SpecificUrn):
3306
+ ENTITY_TYPE: ClassVar[str] = "entityType"
2869
3307
  _URN_PARTS: ClassVar[int] = 1
2870
3308
 
2871
- def __init__(self, id: str, *, _allow_coercion: bool = True) -> None:
3309
+ def __init__(self, id: Union["EntityTypeUrn", str], *, _allow_coercion: bool = True) -> None:
2872
3310
  if _allow_coercion:
2873
3311
  # Field coercion logic (if any is required).
2874
- id = UrnEncoder.encode_string(id)
3312
+ if isinstance(id, str):
3313
+ if id.startswith('urn:li:'):
3314
+ try:
3315
+ id = EntityTypeUrn.from_string(id)
3316
+ except InvalidUrnError:
3317
+ raise InvalidUrnError(f'Expecting a EntityTypeUrn but got {id}')
3318
+ else:
3319
+ id = UrnEncoder.encode_string(id)
2875
3320
 
2876
3321
  # Validation logic.
2877
3322
  if not id:
2878
- raise InvalidUrnError("BusinessAttributeUrn id cannot be empty")
3323
+ raise InvalidUrnError("EntityTypeUrn id cannot be empty")
3324
+ if isinstance(id, EntityTypeUrn):
3325
+ id = id.id
3326
+ elif isinstance(id, Urn):
3327
+ raise InvalidUrnError(f'Expecting a EntityTypeUrn but got {id}')
2879
3328
  if UrnEncoder.contains_reserved_char(id):
2880
- raise InvalidUrnError(f'BusinessAttributeUrn id contains reserved characters')
3329
+ raise InvalidUrnError(f'EntityTypeUrn id contains reserved characters')
2881
3330
 
2882
3331
  super().__init__(self.ENTITY_TYPE, [id])
2883
3332
 
2884
3333
  @classmethod
2885
- def _parse_ids(cls, entity_ids: List[str]) -> "BusinessAttributeUrn":
3334
+ def _parse_ids(cls, entity_ids: List[str]) -> "EntityTypeUrn":
2886
3335
  if len(entity_ids) != cls._URN_PARTS:
2887
- raise InvalidUrnError(f"BusinessAttributeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
3336
+ raise InvalidUrnError(f"EntityTypeUrn should have {cls._URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
2888
3337
  return cls(id=entity_ids[0], _allow_coercion=False)
2889
3338
 
2890
3339
  @classmethod
2891
- def underlying_key_aspect_type(cls) -> Type["BusinessAttributeKeyClass"]:
2892
- from datahub.metadata.schema_classes import BusinessAttributeKeyClass
3340
+ def underlying_key_aspect_type(cls) -> Type["EntityTypeKeyClass"]:
3341
+ from datahub.metadata.schema_classes import EntityTypeKeyClass
2893
3342
 
2894
- return BusinessAttributeKeyClass
3343
+ return EntityTypeKeyClass
2895
3344
 
2896
- def to_key_aspect(self) -> "BusinessAttributeKeyClass":
2897
- from datahub.metadata.schema_classes import BusinessAttributeKeyClass
3345
+ def to_key_aspect(self) -> "EntityTypeKeyClass":
3346
+ from datahub.metadata.schema_classes import EntityTypeKeyClass
2898
3347
 
2899
- return BusinessAttributeKeyClass(id=self.id)
3348
+ return EntityTypeKeyClass(id=self.id)
2900
3349
 
2901
3350
  @classmethod
2902
- def from_key_aspect(cls, key_aspect: "BusinessAttributeKeyClass") -> "BusinessAttributeUrn":
3351
+ def from_key_aspect(cls, key_aspect: "EntityTypeKeyClass") -> "EntityTypeUrn":
2903
3352
  return cls(id=key_aspect.id)
2904
3353
 
2905
3354
  @property