feat(ingestion) Adding vertexAI ingestion source (v1 - model group and model)

datahub-web-react/src/app/ingest/source/builder/sources.json

@@ -333,5 +333,12 @@
         "description": "Import Nodes and Relationships from Neo4j.",
         "docsUrl": "https://datahubproject.io/docs/generated/ingestion/sources/neo4j/",
         "recipe": "source:\n    type: 'neo4j'\n    config:\n        uri: 'neo4j+ssc://host:7687'\n        username: 'neo4j'\n        password: 'password'\n        env: 'PROD'\n\nsink:\n  type: \"datahub-rest\"\n  config:\n    server: 'http://localhost:8080'"
+    },
+    {
+        "urn": "urn:li:dataPlatform:vertexai",
+        "name": "vertexai",
+        "displayName": "VertexAI",
+        "docsUrl": "https://datahubproject.io/docs/generated/ingestion/sources/vertexai/",
+        "recipe": "source:\n  type: vertexai\n  config:\n    project_id: # you GCP project ID \n     region: # region where your GCP project resides  \n   # Credentials\n        # Add GCP credentials"
     }
 ]

metadata-ingestion/docs/sources/vertexai/vertexai_pre.md

@@ -0,0 +1,48 @@
+Ingesting metadata from VertexAI requires using the **Vertex AI** module.
+
+#### Prerequisites
+Please refer to the [Vertex AI documentation](https://cloud.google.com/vertex-ai/docs) for basic information on Vertex AI.
+
+#### Credentials to access to GCP
+Please read the section to understand how to set up application default Credentials to GCP [GCP docs](https://cloud.google.com/docs/authentication/provide-credentials-adc#how-to).
+
+#### Create a service account and assign roles
+
+1. Setup a ServiceAccount as per [GCP docs](https://cloud.google.com/iam/docs/creating-managing-service-accounts#iam-service-accounts-create-console) and assign the previously created role to this service account.
+2. Download a service account JSON keyfile. 
+- Example credential file:
+
+```json
+{
+   "type": "service_account",
+   "project_id": "project-id-1234567",
+   "private_key_id": "d0121d0000882411234e11166c6aaa23ed5d74e0",
+   "private_key": "-----BEGIN PRIVATE KEY-----\nMIIyourkey\n-----END PRIVATE KEY-----",
+   "client_email": "test@suppproject-id-1234567.iam.gserviceaccount.com",
+   "client_id": "113545814931671546333",
+   "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+   "token_uri": "https://oauth2.googleapis.com/token",
+   "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+   "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test%suppproject-id-1234567.iam.gserviceaccount.com"
+}
+```
+
+3. To provide credentials to the source, you can either:
+
+- Set an environment variable:
+
+   ```sh
+   $ export GOOGLE_APPLICATION_CREDENTIALS="/path/to/keyfile.json"
+   ```
+
+   _or_
+
+- Set credential config in your source based on the credential json file. For example:
+
+   ```yml
+   credential:
+     private_key_id: "d0121d0000882411234e11166c6aaa23ed5d74e0"
+     private_key: "-----BEGIN PRIVATE KEY-----\nMIIyourkey\n-----END PRIVATE KEY-----\n"
+     client_email: "test@suppproject-id-1234567.iam.gserviceaccount.com"
+     client_id: "123456678890"
+   ```

metadata-ingestion/docs/sources/vertexai/vertexai_recipe.yml

@@ -0,0 +1,16 @@
+source:
+  type: vertexai
+  config:
+    project_id: "acryl-poc"
+    region:  "us-west2"
+# Note that GOOGLE_APPLICATION_CREDENTIALS or credential section below is required for authentication.
+#   credential:
+#      private_key: '-----BEGIN PRIVATE KEY-----\\nprivate-key\\n-----END PRIVATE KEY-----\\n'
+#      private_key_id: "project_key_id"
+#      client_email: "client_email"
+#      client_id: "client_id"
+
+sink:
+  type: "datahub-rest"
+  config:
+    server: "http://localhost:8080"

metadata-ingestion/setup.py

@@ -532,6 +532,7 @@
     "sigma": sqlglot_lib | {"requests"},
     "sac": sac,
     "neo4j": {"pandas", "neo4j"},
+    "vertexai": {"google-cloud-aiplatform>=1.80.0"},
 }
 
 # This is mainly used to exclude plugins from the Docker image.
@@ -677,6 +678,7 @@
             "sac",
             "cassandra",
             "neo4j",
+            "vertexai",
         ]
         if plugin
         for dependency in plugins[plugin]
@@ -710,6 +712,7 @@
             "mariadb",
             "redash",
             "vertica",
+            "vertexai"
         ]
         if plugin
         for dependency in plugins[plugin]
@@ -799,6 +802,7 @@
         "sac = datahub.ingestion.source.sac.sac:SACSource",
         "cassandra = datahub.ingestion.source.cassandra.cassandra:CassandraSource",
         "neo4j = datahub.ingestion.source.neo4j.neo4j_source:Neo4jSource",
+        "vertexai = datahub.ingestion.source.vertexai:VertexAISource",
     ],
     "datahub.ingestion.transformer.plugins": [
         "pattern_cleanup_ownership = datahub.ingestion.transformer.pattern_cleanup_ownership:PatternCleanUpOwnership",

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py

@@ -1,8 +1,6 @@
-import json
 import logging
 import os
 import re
-import tempfile
 from datetime import timedelta
 from typing import Any, Dict, List, Optional, Union
 
@@ -17,10 +15,10 @@
     PlatformInstanceConfigMixin,
 )
 from datahub.configuration.validate_field_removal import pydantic_removed_field
-from datahub.configuration.validate_multiline_string import pydantic_multiline_string
 from datahub.ingestion.glossary.classification_mixin import (
     ClassificationSourceConfigMixin,
 )
+from datahub.ingestion.source.common.gcp_credentials_config import GCPCredential
 from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig, SQLFilterConfig
 from datahub.ingestion.source.state.stateful_ingestion_base import (
@@ -107,50 +105,8 @@ class BigQueryUsageConfig(BaseUsageConfig):
     )
 
 
-class BigQueryCredential(ConfigModel):
-    project_id: str = Field(description="Project id to set the credentials")
-    private_key_id: str = Field(description="Private key id")
-    private_key: str = Field(
-        description="Private key in a form of '-----BEGIN PRIVATE KEY-----\\nprivate-key\\n-----END PRIVATE KEY-----\\n'"
-    )
-    client_email: str = Field(description="Client email")
-    client_id: str = Field(description="Client Id")
-    auth_uri: str = Field(
-        default="https://accounts.google.com/o/oauth2/auth",
-        description="Authentication uri",
-    )
-    token_uri: str = Field(
-        default="https://oauth2.googleapis.com/token", description="Token uri"
-    )
-    auth_provider_x509_cert_url: str = Field(
-        default="https://www.googleapis.com/oauth2/v1/certs",
-        description="Auth provider x509 certificate url",
-    )
-    type: str = Field(default="service_account", description="Authentication type")
-    client_x509_cert_url: Optional[str] = Field(
-        default=None,
-        description="If not set it will be default to https://www.googleapis.com/robot/v1/metadata/x509/client_email",
-    )
-
-    _fix_private_key_newlines = pydantic_multiline_string("private_key")
-
-    @root_validator(skip_on_failure=True)
-    def validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("client_x509_cert_url") is None:
-            values["client_x509_cert_url"] = (
-                f"https://www.googleapis.com/robot/v1/metadata/x509/{values['client_email']}"
-            )
-        return values
-
-    def create_credential_temp_file(self) -> str:
-        with tempfile.NamedTemporaryFile(delete=False) as fp:
-            cred_json = json.dumps(self.dict(), indent=4, separators=(",", ": "))
-            fp.write(cred_json.encode())
-            return fp.name
-
-
 class BigQueryConnectionConfig(ConfigModel):
-    credential: Optional[BigQueryCredential] = Field(
+    credential: Optional[GCPCredential] = Field(
         default=None, description="BigQuery credential informations"
     )
 

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

@@ -0,0 +1,53 @@
+import json
+import tempfile
+from typing import Any, Dict, Optional
+
+from pydantic import Field, root_validator
+
+from datahub.configuration import ConfigModel
+from datahub.configuration.validate_multiline_string import pydantic_multiline_string
+
+
+class GCPCredential(ConfigModel):
+    project_id: Optional[str] = Field(description="Project id to set the credentials")
+    private_key_id: str = Field(description="Private key id")
+    private_key: str = Field(
+        description="Private key in a form of '-----BEGIN PRIVATE KEY-----\\nprivate-key\\n-----END PRIVATE KEY-----\\n'"
+    )
+    client_email: str = Field(description="Client email")
+    client_id: str = Field(description="Client Id")
+    auth_uri: str = Field(
+        default="https://accounts.google.com/o/oauth2/auth",
+        description="Authentication uri",
+    )
+    token_uri: str = Field(
+        default="https://oauth2.googleapis.com/token", description="Token uri"
+    )
+    auth_provider_x509_cert_url: str = Field(
+        default="https://www.googleapis.com/oauth2/v1/certs",
+        description="Auth provider x509 certificate url",
+    )
+    type: str = Field(default="service_account", description="Authentication type")
+    client_x509_cert_url: Optional[str] = Field(
+        default=None,
+        description="If not set it will be default to https://www.googleapis.com/robot/v1/metadata/x509/client_email",
+    )
+
+    _fix_private_key_newlines = pydantic_multiline_string("private_key")
+
+    @root_validator(skip_on_failure=True)
+    def validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        if values.get("client_x509_cert_url") is None:
+            values["client_x509_cert_url"] = (
+                f"https://www.googleapis.com/robot/v1/metadata/x509/{values['client_email']}"
+            )
+        return values
+
+    def create_credential_temp_file(self, project_id: Optional[str] = None) -> str:
+        configs = self.dict()
+        if project_id:
+            configs["project_id"] = project_id
+        with tempfile.NamedTemporaryFile(delete=False) as fp:
+            cred_json = json.dumps(configs, indent=4, separators=(",", ": "))
+            fp.write(cred_json.encode())
+            return fp.name