[SNOW-2437420] Add experiment to lineage versioned domains (#3934)

Author: sfc-gh-yuzhao

src/snowflake/snowpark/lineage.py

@@ -135,6 +135,7 @@ class _SnowflakeDomain:
     TABLE = "TABLE"
     MODULE = "MODULE"
     DATASET = "DATASET"
+    EXPERIMENT = "EXPERIMENT"
     VIEW = "VIEW"
     COLUMN = "COLUMN"
     SNOWSERVICE_INSTANCE = "SNOWSERVICE_INSTANCE"
@@ -274,6 +275,7 @@ def __init__(self, session: "snowflake.snowpark.session.Session") -> None:
             _UserDomain.FEATURE_VIEW,
             _UserDomain.MODEL,
             _SnowflakeDomain.DATASET,
+            _SnowflakeDomain.EXPERIMENT,
         }
 
     def _get_lineage(

tests/unit/test_lineage.py

@@ -112,6 +112,17 @@ def test_get_name_and_version():
         in str(exc)
     )
 
+    graph_entity = {
+        _ObjectField.USER_DOMAIN: _SnowflakeDomain.EXPERIMENT,
+        _ObjectField.DB: "db1",
+        _ObjectField.SCHEMA: "schema1",
+        _ObjectField.PROPERTIES: {_ObjectField.PARENT_NAME: "exp_parent"},
+        _ObjectField.NAME: "MY_RUN",
+    }
+    name, version = Lineage(fake_session)._get_name_and_version(graph_entity)
+    assert name == "db1.schema1.exp_parent"
+    assert version == "MY_RUN"
+
 
 def test_get_user_entity():
     fake_session = mock.create_autospec(Session, _session_id=123456)
@@ -213,6 +224,22 @@ def test_get_user_entity():
     assert _ObjectField.TYPE in user_entity
     assert user_entity[_ObjectField.TYPE] == "TABLE"
 
+    # EXPERIMENT should surface name and version
+    graph_entity = {
+        _ObjectField.USER_DOMAIN: _SnowflakeDomain.EXPERIMENT,
+        _ObjectField.DB: "db1",
+        _ObjectField.SCHEMA: "schema1",
+        _ObjectField.PROPERTIES: {_ObjectField.PARENT_NAME: "exp_parent"},
+        _ObjectField.NAME: "MY_RUN",
+        _ObjectField.CREATED_ON: "123455",
+        _ObjectField.STATUS: "Active",
+    }
+
+    user_entity = Lineage(fake_session)._get_user_entity(graph_entity)
+    assert user_entity[_ObjectField.NAME] == "db1.schema1.exp_parent"
+    assert user_entity[_ObjectField.DOMAIN] == _SnowflakeDomain.EXPERIMENT
+    assert user_entity[_ObjectField.VERSION] == "MY_RUN"
+
 
 def test_split_fully_qualified_name():
     test_cases_valid = [