Removing temp credential file creation

Author: ryota-cloud

metadata-ingestion/src/datahub/ingestion/source/common/gcp_credentials_config.py

@@ -51,3 +51,9 @@ def create_credential_temp_file(self, project_id: Optional[str] = None) -> str:
             cred_json = json.dumps(configs, indent=4, separators=(",", ": "))
             fp.write(cred_json.encode())
             return fp.name
+
+    def to_dict(self, project_id: Optional[str] = None) -> Dict[str, str]:
+        configs = self.dict()
+        if project_id:
+            configs["project_id"] = project_id
+        return configs

metadata-ingestion/src/datahub/ingestion/source/vertexai/config.py

@@ -1,6 +1,6 @@
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
-from pydantic import Field, PrivateAttr
+from pydantic import Field
 
 from datahub.configuration.source_common import EnvConfigMixin
 from datahub.ingestion.source.common.gcp_credentials_config import GCPCredential
@@ -23,11 +23,10 @@ class VertexAIConfig(EnvConfigMixin):
         description=("VertexUI URI"),
     )
 
-    _credentials_path: Optional[str] = PrivateAttr(None)
-
     def __init__(self, **data: Any):
         super().__init__(**data)
+
+    def get_credentials(self) -> Optional[Dict[str, str]]:
         if self.credential:
-            self._credentials_path = self.credential.create_credential_temp_file(
-                project_id=self.project_id
-            )
+            return self.credential.to_dict(self.project_id)
+        return None

metadata-ingestion/src/datahub/ingestion/source/vertexai/vertexai.py

@@ -113,11 +113,10 @@ def __init__(self, ctx: PipelineContext, config: VertexAIConfig):
         self.config = config
         self.report = SourceReport()
 
+        creds = self.config.get_credentials()
         credentials = (
-            service_account.Credentials.from_service_account_file(
-                self.config._credentials_path
-            )
-            if self.config.credential
+            service_account.Credentials.from_service_account_info(**creds)
+            if creds
             else None
         )
 
@@ -294,6 +293,16 @@ def _gen_run_execution(
                     upstreamInstances=[self._make_experiment_run_urn(exp, run)],
                     parentInstance=self._make_experiment_run_urn(exp, run),
                 ),
+                (
+                    DataProcessInstanceInputClass(
+                        inputs=[],
+                        inputEdges=[
+                            EdgeClass(
+                                destinationUrn=self._make_experiment_run_urn(exp, run)
+                            ),
+                        ],
+                    )
+                ),
             ],
         )
 
@@ -342,7 +351,9 @@ def _gen_experiment_run_mcps(
                 (
                     DataProcessInstanceRunEventClass(
                         status=DataProcessRunStatusClass.COMPLETE,
-                        timestampMillis=created_time if created_time else 0, # None is not allowed, 0 as default value
+                        timestampMillis=created_time
+                        if created_time
+                        else 0,  # None is not allowed, 0 as default value
                         result=DataProcessInstanceRunResultClass(
                             type=run_result_type,
                             nativeResultType=self.platform,

metadata-ingestion/tests/unit/test_vertexai_source.py

@@ -1,5 +1,4 @@
 import contextlib
-import json
 from typing import List
 from unittest.mock import patch
 
@@ -110,7 +109,6 @@ def test_get_ml_model_properties_mcps(source: VertexAISource) -> None:
 
     # Run _gen_ml_model_mcps
     actual_mcps = list(source._gen_ml_model_mcps(model_meta))
-
     actual_urns = [mcp.entityUrn for mcp in actual_mcps]
     expected_urns = [
         source._make_ml_model_urn(
@@ -362,31 +360,29 @@ def test_vertexai_config_init():
     assert config.credential.client_id == "test-client-id"
     assert config.credential.auth_uri == "https://accounts.google.com/o/oauth2/auth"
     assert config.credential.token_uri == "https://oauth2.googleapis.com/token"
+    assert config.credential.auth_provider_x509_cert_url == "service_account"
+
+    parsed_conf = config.get_credentials()
+    assert parsed_conf is not None
+    assert parsed_conf.get("project_id") == config_data["project_id"]
+    assert "credential" in config_data
+    assert parsed_conf.get("private_key_id", "") == "test-key-id"
+    assert (
+        parsed_conf.get("private_key", "")
+        == "-----BEGIN PRIVATE KEY-----\ntest-private-key\n-----END PRIVATE KEY-----\n"
+    )
+    assert (
+        parsed_conf.get("client_email")
+        == "[email protected]"
+    )
+    assert parsed_conf.get("client_id") == "test-client-id"
+    assert parsed_conf.get("auth_uri") == "https://accounts.google.com/o/oauth2/auth"
+    assert parsed_conf.get("token_uri") == "https://oauth2.googleapis.com/token"
     assert (
-        config.credential.auth_provider_x509_cert_url
+        parsed_conf.get("auth_provider_x509_cert_url")
         == "https://www.googleapis.com/oauth2/v1/certs"
     )
-
-    assert config._credentials_path is not None
-    with open(config._credentials_path, "r") as file:
-        content = json.loads(file.read())
-        assert content["project_id"] == "test-project"
-        assert content["private_key_id"] == "test-key-id"
-        assert content["private_key_id"] == "test-key-id"
-        assert (
-            content["private_key"]
-            == "-----BEGIN PRIVATE KEY-----\ntest-private-key\n-----END PRIVATE KEY-----\n"
-        )
-        assert (
-            content["client_email"] == "[email protected]"
-        )
-        assert content["client_id"] == "test-client-id"
-        assert content["auth_uri"] == "https://accounts.google.com/o/oauth2/auth"
-        assert content["token_uri"] == "https://oauth2.googleapis.com/token"
-        assert (
-            content["auth_provider_x509_cert_url"]
-            == "https://www.googleapis.com/oauth2/v1/certs"
-        )
+    assert parsed_conf.get("type") == "service_account"
 
 
 def test_get_input_dataset_mcps(source: VertexAISource) -> None: