acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/openapi.py

@@ -2,13 +2,14 @@ import logging
 import time
 import warnings
 from abc import ABC
-from typing import Dict, Iterable, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from pydantic import validator
 from pydantic.fields import Field
 
 from datahub.configuration.common import ConfigModel
 from datahub.emitter.mce_builder import make_tag_urn
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -20,6 +21,7 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import DatasetSubTypes
 from datahub.ingestion.source.openapi_parser import (
     clean_url,
     compose_url_attr,
@@ -32,14 +34,13 @@ from datahub.ingestion.source.openapi_parser import (
     set_metadata,
     try_guessing,
 )
-from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
-from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
     AuditStampClass,
     DatasetPropertiesClass,
     GlobalTagsClass,
     InstitutionalMemoryClass,
     InstitutionalMemoryMetadataClass,
+    SubTypesClass,
     TagAssociationClass,
 )
 
@@ -81,6 +82,9 @@ class OpenApiConfig(ConfigModel):
     get_token: dict = Field(
         default={}, description="Retrieving a token from the endpoint."
     )
+    verify_ssl: bool = Field(
+        default=True, description="Enable SSL certificate verification"
+    )
 
     @validator("bearer_token", always=True)
     def ensure_only_one_token(
@@ -101,7 +105,7 @@ class OpenApiConfig(ConfigModel):
                 # details there once, and then use that session for all requests.
                 self.token = f"Bearer {self.bearer_token}"
             else:
-                assert "url_complement" in self.get_token.keys(), (
+                assert "url_complement" in self.get_token, (
                     "When 'request_type' is set to 'get', an url_complement is needed for the request."
                 )
                 if self.get_token["request_type"] == "get":
@@ -128,12 +132,14 @@ class OpenApiConfig(ConfigModel):
                     tok_url=url4req,
                     method=self.get_token["request_type"],
                     proxies=self.proxies,
+                    verify_ssl=self.verify_ssl,
                 )
             sw_dict = get_swag_json(
                 self.url,
                 token=self.token,
                 swagger_file=self.swagger_file,
                 proxies=self.proxies,
+                verify_ssl=self.verify_ssl,
             )  # load the swagger file
 
         else:  # using basic auth for accessing endpoints
@@ -143,6 +149,7 @@ class OpenApiConfig(ConfigModel):
                 password=self.password,
                 swagger_file=self.swagger_file,
                 proxies=self.proxies,
+                verify_ssl=self.verify_ssl,
             )
         return sw_dict
 
@@ -222,8 +229,9 @@ class APISource(Source, ABC):
 
     def init_dataset(
         self, endpoint_k: str, endpoint_dets: dict
-    ) -> Tuple[DatasetSnapshot, str]:
+    ) -> Tuple[str, str, List[MetadataWorkUnit]]:
         config = self.config
+        workunits = []
 
         dataset_name = endpoint_k[1:].replace("/", ".")
 
@@ -233,22 +241,27 @@ class APISource(Source, ABC):
         else:
             dataset_name = "root"
 
-        dataset_snapshot = DatasetSnapshot(
-            urn=f"urn:li:dataset:(urn:li:dataPlatform:{self.platform},{config.name}.{dataset_name},PROD)",
-            aspects=[],
-        )
+        dataset_urn = f"urn:li:dataset:(urn:li:dataPlatform:{self.platform},{config.name}.{dataset_name},PROD)"
 
-        # adding description
-        dataset_properties = DatasetPropertiesClass(
+        # Create dataset properties aspect
+        properties = DatasetPropertiesClass(
             description=endpoint_dets["description"], customProperties={}
         )
-        dataset_snapshot.aspects.append(dataset_properties)
+        wu = MetadataWorkUnit(
+            id=dataset_name,
+            mcp=MetadataChangeProposalWrapper(entityUrn=dataset_urn, aspect=properties),
+        )
+        workunits.append(wu)
 
-        # adding tags
+        # Create tags aspect
         tags_str = [make_tag_urn(t) for t in endpoint_dets["tags"]]
         tags_tac = [TagAssociationClass(t) for t in tags_str]
         gtc = GlobalTagsClass(tags_tac)
-        dataset_snapshot.aspects.append(gtc)
+        wu = MetadataWorkUnit(
+            id=f"{dataset_name}-tags",
+            mcp=MetadataChangeProposalWrapper(entityUrn=dataset_urn, aspect=gtc),
+        )
+        workunits.append(wu)
 
         # the link will appear in the "documentation"
         link_url = clean_url(config.url + self.url_basepath + endpoint_k)
@@ -260,17 +273,25 @@ class APISource(Source, ABC):
             url=link_url, description=link_description, createStamp=creation
         )
         inst_memory = InstitutionalMemoryClass([link_metadata])
-        dataset_snapshot.aspects.append(inst_memory)
+        wu = MetadataWorkUnit(
+            id=f"{dataset_name}-docs",
+            mcp=MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn, aspect=inst_memory
+            ),
+        )
+        workunits.append(wu)
 
-        return dataset_snapshot, dataset_name
+        # Create subtype aspect
+        sub_types = SubTypesClass(typeNames=[DatasetSubTypes.API_ENDPOINT])
+        wu = MetadataWorkUnit(
+            id=f"{dataset_name}-subtype",
+            mcp=MetadataChangeProposalWrapper(entityUrn=dataset_urn, aspect=sub_types),
+        )
+        workunits.append(wu)
 
-    def build_wu(
-        self, dataset_snapshot: DatasetSnapshot, dataset_name: str
-    ) -> ApiWorkUnit:
-        mce = MetadataChangeEvent(proposedSnapshot=dataset_snapshot)
-        return ApiWorkUnit(id=dataset_name, mce=mce)
+        return dataset_name, dataset_urn, workunits
 
-    def get_workunits_internal(self) -> Iterable[ApiWorkUnit]:
+    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         config = self.config
 
         sw_dict = self.config.get_swagger()
@@ -294,17 +315,25 @@ class APISource(Source, ABC):
             if endpoint_k in config.ignore_endpoints:
                 continue
 
-            dataset_snapshot, dataset_name = self.init_dataset(
+            # Initialize dataset and get common aspects
+            dataset_name, dataset_urn, workunits = self.init_dataset(
                 endpoint_k, endpoint_dets
             )
+            for wu in workunits:
+                yield wu
 
-            # adding dataset fields
-            if "data" in endpoint_dets.keys():
+            # Handle schema metadata if available
+            if "data" in endpoint_dets:
                 # we are lucky! data is defined in the swagger for this endpoint
                 schema_metadata = set_metadata(dataset_name, endpoint_dets["data"])
-                dataset_snapshot.aspects.append(schema_metadata)
-                yield self.build_wu(dataset_snapshot, dataset_name)
-            elif endpoint_dets["method"] != "get":
+                wu = MetadataWorkUnit(
+                    id=f"{dataset_name}-schema",
+                    mcp=MetadataChangeProposalWrapper(
+                        entityUrn=dataset_urn, aspect=schema_metadata
+                    ),
+                )
+                yield wu
+            elif endpoint_dets["method"] != "GET":
                 self.report.report_warning(
                     title="Failed to Extract Endpoint Metadata",
                     message=f"No example provided for {endpoint_dets['method']}",
@@ -320,6 +349,7 @@ class APISource(Source, ABC):
                         tot_url,
                         token=config.token,
                         proxies=config.proxies,
+                        verify_ssl=config.verify_ssl,
                     )
                 else:
                     response = request_call(
@@ -327,6 +357,7 @@ class APISource(Source, ABC):
                         username=config.username,
                         password=config.password,
                         proxies=config.proxies,
+                        verify_ssl=config.verify_ssl,
                     )
                 if response.status_code == 200:
                     fields2add, root_dataset_samples[dataset_name] = extract_fields(
@@ -338,13 +369,17 @@ class APISource(Source, ABC):
                             context=f"Endpoint Type: {endpoint_k}, Name: {dataset_name}",
                         )
                     schema_metadata = set_metadata(dataset_name, fields2add)
-                    dataset_snapshot.aspects.append(schema_metadata)
-
-                    yield self.build_wu(dataset_snapshot, dataset_name)
+                    wu = MetadataWorkUnit(
+                        id=f"{dataset_name}-schema",
+                        mcp=MetadataChangeProposalWrapper(
+                            entityUrn=dataset_urn, aspect=schema_metadata
+                        ),
+                    )
+                    yield wu
                 else:
                     self.report_bad_responses(response.status_code, type=endpoint_k)
             else:
-                if endpoint_k not in config.forced_examples.keys():
+                if endpoint_k not in config.forced_examples:
                     # start guessing...
                     url_guess = try_guessing(endpoint_k, root_dataset_samples)
                     tot_url = clean_url(config.url + self.url_basepath + url_guess)
@@ -353,6 +388,7 @@ class APISource(Source, ABC):
                             tot_url,
                             token=config.token,
                             proxies=config.proxies,
+                            verify_ssl=config.verify_ssl,
                         )
                     else:
                         response = request_call(
@@ -360,6 +396,7 @@ class APISource(Source, ABC):
                             username=config.username,
                             password=config.password,
                             proxies=config.proxies,
+                            verify_ssl=config.verify_ssl,
                         )
                     if response.status_code == 200:
                         fields2add, _ = extract_fields(response, dataset_name)
@@ -369,9 +406,13 @@ class APISource(Source, ABC):
                                 context=f"Endpoint Type: {endpoint_k}, Name: {dataset_name}",
                             )
                         schema_metadata = set_metadata(dataset_name, fields2add)
-                        dataset_snapshot.aspects.append(schema_metadata)
-
-                        yield self.build_wu(dataset_snapshot, dataset_name)
+                        wu = MetadataWorkUnit(
+                            id=f"{dataset_name}-schema",
+                            mcp=MetadataChangeProposalWrapper(
+                                entityUrn=dataset_urn, aspect=schema_metadata
+                            ),
+                        )
+                        yield wu
                     else:
                         self.report_bad_responses(response.status_code, type=endpoint_k)
                 else:
@@ -384,6 +425,7 @@ class APISource(Source, ABC):
                             tot_url,
                             token=config.token,
                             proxies=config.proxies,
+                            verify_ssl=config.verify_ssl,
                         )
                     else:
                         response = request_call(
@@ -391,6 +433,7 @@ class APISource(Source, ABC):
                             username=config.username,
                             password=config.password,
                             proxies=config.proxies,
+                            verify_ssl=config.verify_ssl,
                         )
                     if response.status_code == 200:
                         fields2add, _ = extract_fields(response, dataset_name)
@@ -400,9 +443,13 @@ class APISource(Source, ABC):
                                 context=f"Endpoint Type: {endpoint_k}, Name: {dataset_name}",
                             )
                         schema_metadata = set_metadata(dataset_name, fields2add)
-                        dataset_snapshot.aspects.append(schema_metadata)
-
-                        yield self.build_wu(dataset_snapshot, dataset_name)
+                        wu = MetadataWorkUnit(
+                            id=f"{dataset_name}-schema",
+                            mcp=MetadataChangeProposalWrapper(
+                                entityUrn=dataset_urn, aspect=schema_metadata
+                            ),
+                        )
+                        yield wu
                     else:
                         self.report_bad_responses(response.status_code, type=endpoint_k)
 

datahub/ingestion/source/openapi_parser.py

@@ -59,17 +59,21 @@ def request_call(
     username: Optional[str] = None,
     password: Optional[str] = None,
     proxies: Optional[dict] = None,
+    verify_ssl: bool = True,
 ) -> requests.Response:
     headers = {"accept": "application/json"}
     if username is not None and password is not None:
         return requests.get(
-            url, headers=headers, auth=HTTPBasicAuth(username, password)
+            url,
+            headers=headers,
+            auth=HTTPBasicAuth(username, password),
+            verify=verify_ssl,
         )
     elif token is not None:
         headers["Authorization"] = f"{token}"
-        return requests.get(url, proxies=proxies, headers=headers)
+        return requests.get(url, proxies=proxies, headers=headers, verify=verify_ssl)
     else:
-        return requests.get(url, headers=headers)
+        return requests.get(url, headers=headers, verify=verify_ssl)
 
 
 def get_swag_json(
@@ -79,10 +83,16 @@ def get_swag_json(
     password: Optional[str] = None,
     swagger_file: str = "",
     proxies: Optional[dict] = None,
+    verify_ssl: bool = True,
 ) -> Dict:
     tot_url = url + swagger_file
     response = request_call(
-        url=tot_url, token=token, username=username, password=password, proxies=proxies
+        url=tot_url,
+        token=token,
+        username=username,
+        password=password,
+        proxies=proxies,
+        verify_ssl=verify_ssl,
     )
 
     if response.status_code != 200:
@@ -127,37 +137,45 @@ def get_endpoints(sw_dict: dict) -> dict:
     check_sw_version(sw_dict)
 
     for p_k, p_o in sw_dict["paths"].items():
-        method = list(p_o)[0]
-        if "200" in p_o[method]["responses"].keys():
-            base_res = p_o[method]["responses"]["200"]
-        elif 200 in p_o[method]["responses"].keys():
-            # if you read a plain yml file the 200 will be an integer
-            base_res = p_o[method]["responses"][200]
-        else:
-            # the endpoint does not have a 200 response
-            continue
-
-        if "description" in p_o[method].keys():
-            desc = p_o[method]["description"]
-        elif "summary" in p_o[method].keys():
-            desc = p_o[method]["summary"]
-        else:  # still testing
-            desc = ""
-
-        try:
-            tags = p_o[method]["tags"]
-        except KeyError:
-            tags = []
-
-        url_details[p_k] = {"description": desc, "tags": tags, "method": method}
-
-        example_data = check_for_api_example_data(base_res, p_k)
-        if example_data:
-            url_details[p_k]["data"] = example_data
-
-        # checking whether there are defined parameters to execute the call...
-        if "parameters" in p_o[method].keys():
-            url_details[p_k]["parameters"] = p_o[method]["parameters"]
+        for method, method_spec in p_o.items():
+            # skip non-method keys like "parameters"
+            if method.lower() not in [
+                "get",
+                "post",
+                "put",
+                "delete",
+                "patch",
+                "options",
+                "head",
+            ]:
+                continue
+
+            responses = method_spec.get("responses", {})
+            base_res = responses.get("200") or responses.get(200)
+            if not base_res:
+                # if there is no 200 response, we skip this method
+                continue
+
+            # if the description is not present, we will use the summary
+            # if both are not present, we will use an empty string
+            desc = method_spec.get("description") or method_spec.get("summary", "")
+
+            # if the tags are not present, we will use an empty list
+            tags = method_spec.get("tags", [])
+
+            url_details[p_k] = {
+                "description": desc,
+                "tags": tags,
+                "method": method.upper(),
+            }
+
+            example_data = check_for_api_example_data(base_res, p_k)
+            if example_data:
+                url_details[p_k]["data"] = example_data
+
+            # checking whether there are defined parameters to execute the call...
+            if "parameters" in p_o[method]:
+                url_details[p_k]["parameters"] = p_o[method]["parameters"]
 
     return dict(sorted(url_details.items()))
 
@@ -169,7 +187,7 @@ def check_for_api_example_data(base_res: dict, key: str) -> dict:
     data = {}
     if "content" in base_res:
         res_cont = base_res["content"]
-        if "application/json" in res_cont.keys():
+        if "application/json" in res_cont:
             ex_field = None
             if "example" in res_cont["application/json"]:
                 ex_field = "example"
@@ -186,7 +204,7 @@ def check_for_api_example_data(base_res: dict, key: str) -> dict:
                 logger.warning(
                     f"Field in swagger file does not give consistent data --- {key}"
                 )
-        elif "text/csv" in res_cont.keys():
+        elif "text/csv" in res_cont:
             data = res_cont["text/csv"]["schema"]
     elif "examples" in base_res:
         data = base_res["examples"]["application/json"]
@@ -239,7 +257,7 @@ def guessing_url_name(url: str, examples: dict) -> str:
 
     # substituting the parameter's name w the value
     for name, clean_name in zip(needed_n, cleaned_needed_n):
-        if clean_name in examples[ex2use].keys():
+        if clean_name in examples[ex2use]:
             guessed_url = re.sub(name, str(examples[ex2use][clean_name]), guessed_url)
 
     return guessed_url
@@ -358,6 +376,7 @@ def get_tok(
     tok_url: str = "",
     method: str = "post",
     proxies: Optional[dict] = None,
+    verify_ssl: bool = True,
 ) -> str:
     """
     Trying to post username/password to get auth.
@@ -368,7 +387,7 @@ def get_tok(
         # this will make a POST call with username and password
         data = {"username": username, "password": password, "maxDuration": True}
         # url2post = url + "api/authenticate/"
-        response = requests.post(url4req, proxies=proxies, json=data)
+        response = requests.post(url4req, proxies=proxies, json=data, verify=verify_ssl)
         if response.status_code == 200:
             cont = json.loads(response.content)
             if "token" in cont:  # other authentication scheme
@@ -377,7 +396,7 @@ def get_tok(
                 token = f"Bearer {cont['tokens']['access']}"
     elif method == "get":
         # this will make a GET call with username and password
-        response = requests.get(url4req)
+        response = requests.get(url4req, verify=verify_ssl)
         if response.status_code == 200:
             cont = json.loads(response.content)
             token = cont["token"]