feat(ldap): add attrs_list parameter and test

datahub-project · Jul 27, 2022 · 6ab88c3 · 6ab88c3
1 parent 477d118
commit 6ab88c3
Show file tree

Hide file tree

Showing 5 changed files with 269 additions and 0 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/ldap.py b/metadata-ingestion/src/datahub/ingestion/source/ldap.py
@@ -97,6 +97,7 @@ class LDAPSourceConfig(ConfigModel):
     # Extraction configuration.
     base_dn: str = Field(description="LDAP DN.")
     filter: str = Field(default="(objectClass=*)", description="LDAP extractor filter.")
+    attrs_list: List[str] = Field(default=None, description="Retrieved attributes list")
 
     # If set to true, any users without first and last names will be dropped.
     drop_missing_first_last_name: bool = Field(
@@ -204,6 +205,7 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
                     self.config.base_dn,
                     ldap.SCOPE_SUBTREE,
                     self.config.filter,
+                    self.config.attrs_list,
                     serverctrls=[self.lc],
                 )
                 _rtype, rdata, _rmsgid, serverctrls = self.ldap_client.result3(msgid)
@@ -231,6 +233,7 @@ def get_workunits(self) -> Iterable[MetadataWorkUnit]:
                 elif (
                     b"posixGroup" in attrs["objectClass"]
                     or b"organizationalUnit" in attrs["objectClass"]
+                    or b"groupOfNames" in attrs["objectClass"]
                     or b"group" in attrs["objectClass"]
                 ):
                     yield from self.handle_group(dn, attrs)

diff --git a/metadata-ingestion/tests/integration/ldap/ldap_mces_golden.json b/metadata-ingestion/tests/integration/ldap/ldap_mces_golden.json
@@ -151,5 +151,93 @@
         "runId": "ldap-test",
         "properties": null
     }
+},
+{
+    "auditHeader": null,
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+            "urn": "urn:li:corpuser:hbevan",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                        "active": true,
+                        "displayName": "Hester Bevan",
+                        "email": "hbevan",
+                        "title": null,
+                        "managerUrn": null,
+                        "departmentId": null,
+                        "departmentName": null,
+                        "firstName": "Hester",
+                        "lastName": "Bevan",
+                        "fullName": "Hester Bevan",
+                        "countryCode": null
+                    }
+                }
+            ]
+        }
+    },
+    "proposedDelta": null,
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "ldap-test",
+        "properties": null
+    }
+},
+{
+    "auditHeader": null,
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+            "urn": "urn:li:corpuser:ehaas",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                        "active": true,
+                        "displayName": "Evalyn Haas",
+                        "email": "ehaas",
+                        "title": null,
+                        "managerUrn": null,
+                        "departmentId": null,
+                        "departmentName": null,
+                        "firstName": "Evalyn",
+                        "lastName": "Haas",
+                        "fullName": "Evalyn Haas",
+                        "countryCode": null
+                    }
+                }
+            ]
+        }
+    },
+    "proposedDelta": null,
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "ldap-test",
+        "properties": null
+    }
+},
+{
+    "auditHeader": null,
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.CorpGroupSnapshot": {
+            "urn": "urn:li:corpGroup:HR Department",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.identity.CorpGroupInfo": {
+                        "displayName": null,
+                        "email": "HR Department",
+                        "admins": [],
+                        "members": [],
+                        "groups": [],
+                        "description": null
+                    }
+                }
+            ]
+        }
+    },
+    "proposedDelta": null,
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "ldap-test",
+        "properties": null
+    }
 }
 ]
diff --git a/metadata-ingestion/tests/integration/ldap/ldap_memberof_mces_golden.json b/metadata-ingestion/tests/integration/ldap/ldap_memberof_mces_golden.json
@@ -0,0 +1,82 @@
+[
+{
+    "auditHeader": null,
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+            "urn": "urn:li:corpuser:hbevan",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                        "active": true,
+                        "displayName": "Hester Bevan",
+                        "email": "hbevan",
+                        "title": null,
+                        "managerUrn": null,
+                        "departmentId": null,
+                        "departmentName": null,
+                        "firstName": "Hester",
+                        "lastName": "Bevan",
+                        "fullName": "Hester Bevan",
+                        "countryCode": null
+                    }
+                },
+                {
+                    "com.linkedin.pegasus2avro.identity.GroupMembership": {
+                        "groups": [
+                            "urn:li:corpGroup:HR Department"
+                        ]
+                    }
+                }
+            ]
+        }
+    },
+    "proposedDelta": null,
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "ldap-test",
+        "registryName": null,
+        "registryVersion": null,
+        "properties": null
+    }
+},
+{
+    "auditHeader": null,
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.CorpUserSnapshot": {
+            "urn": "urn:li:corpuser:ehaas",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.identity.CorpUserInfo": {
+                        "active": true,
+                        "displayName": "Evalyn Haas",
+                        "email": "ehaas",
+                        "title": null,
+                        "managerUrn": null,
+                        "departmentId": null,
+                        "departmentName": null,
+                        "firstName": "Evalyn",
+                        "lastName": "Haas",
+                        "fullName": "Evalyn Haas",
+                        "countryCode": null
+                    }
+                },
+                {
+                    "com.linkedin.pegasus2avro.identity.GroupMembership": {
+                        "groups": [
+                            "urn:li:corpGroup:HR Department"
+                        ]
+                    }
+                }
+            ]
+        }
+    },
+    "proposedDelta": null,
+    "systemMetadata": {
+        "lastObserved": 1615443388097,
+        "runId": "ldap-test",
+        "registryName": null,
+        "registryVersion": null,
+        "properties": null
+    }
+}
+]
diff --git a/metadata-ingestion/tests/integration/ldap/setup/custom/sample.ldif b/metadata-ingestion/tests/integration/ldap/setup/custom/sample.ldif
@@ -9,6 +9,19 @@
 
 version: 1
 
+# Enable memberOf overlay
+dn: olcOverlay=memberof,olcDatabase={1}mdb,cn=config
+objectClass: olcMemberOf
+objectClass: olcOverlayConfig
+objectClass: olcConfig
+objectClass: top
+olcOverlay: memberof
+olcMemberOfDangling: ignore
+olcMemberOfRefInt: TRUE
+olcMemberOfGroupOC: groupOfNames
+olcMemberOfMemberAD: member
+olcMemberOfMemberOfAD: memberOf
+
 # Entry 1: dc=example,dc=org
 # Note: this is commented out because the containers bootstrap scripts already
 # handle this for us.
@@ -113,3 +126,39 @@ dn: ou=Sales Department,dc=example,dc=org
 objectclass: organizationalUnit
 objectclass: top
 ou: Sales Department
+
+# Entry 10: cn=Hester Bevan,ou=people,dc=example,dc=org
+dn: cn=Hester Bevan,ou=people,dc=example,dc=org
+cn: Hester Bevan
+gidnumber: 500
+givenname: Hester
+homedirectory: /home/users/hbevan
+objectclass: inetOrgPerson
+objectclass: posixAccount
+objectclass: top
+sn: Bevan
+uid: hbevan
+uidnumber: 1004
+userpassword: {MD5}4QrcOUm6Wau+VuBX8g+IPg==
+
+# Entry 11: cn=Evalyn Haas,ou=people,dc=example,dc=org
+dn: cn=Evalyn Haas,ou=people,dc=example,dc=org
+cn: Evalyn Haas
+gidnumber: 500
+givenname: Evalyn
+homedirectory: /home/users/ehaas
+objectclass: inetOrgPerson
+objectclass: posixAccount
+objectclass: top
+sn: Haas
+uid: ehaas
+uidnumber: 1005
+userpassword: {MD5}4QrcOUm6Wau+VuBX8g+IPg==
+
+# Entry 12: cn=HR Department,ou=groups,dc=example,dc=org
+dn: cn=HR Department,dc=example,dc=org
+cn: HR Department
+objectclass: groupOfNames
+objectclass: top
+member: cn=Hester Bevan,ou=people,dc=example,dc=org
+member: cn=Evalyn Haas,ou=people,dc=example,dc=org
diff --git a/metadata-ingestion/tests/integration/ldap/test_ldap.py b/metadata-ingestion/tests/integration/ldap/test_ldap.py
@@ -50,3 +50,50 @@ def test_ldap_ingest(docker_compose_runner, pytestconfig, tmp_path, mock_time):
             output_path=tmp_path / "ldap_mces.json",
             golden_path=test_resources_dir / "ldap_mces_golden.json",
         )
+
+
+@pytest.mark.integration
+def test_ldap_memberof_ingest(docker_compose_runner, pytestconfig, tmp_path, mock_time):
+    test_resources_dir = pytestconfig.rootpath / "tests/integration/ldap"
+
+    with docker_compose_runner(
+        test_resources_dir / "docker-compose.yml", "ldap"
+    ) as docker_services:
+        # The openldap container loads the sample data after exposing the port publicly. As such,
+        # we must wait a little bit extra to ensure that the sample data is loaded.
+        wait_for_port(docker_services, "openldap", 389)
+        time.sleep(5)
+
+        pipeline = Pipeline.create(
+            {
+                "run_id": "ldap-test",
+                "source": {
+                    "type": "ldap",
+                    "config": {
+                        "ldap_server": "ldap://localhost",
+                        "ldap_user": "cn=admin,dc=example,dc=org",
+                        "ldap_password": "admin",
+                        "base_dn": "dc=example,dc=org",
+                        "filter": "(memberOf=cn=HR Department,dc=example,dc=org)",
+                        "attrs_list": ["+", "*"],
+                        "group_attrs_map": {
+                            "members": "member",
+                        },
+                    },
+                },
+                "sink": {
+                    "type": "file",
+                    "config": {
+                        "filename": f"{tmp_path}/ldap_memberof_mces.json",
+                    },
+                },
+            }
+        )
+        pipeline.run()
+        pipeline.raise_from_status()
+
+        mce_helpers.check_golden_file(
+            pytestconfig,
+            output_path=tmp_path / "ldap_memberof_mces.json",
+            golden_path=test_resources_dir / "ldap_memberof_mces_golden.json",
+        )