From 5aba41ef58516474d53beddf5dbc48b63088bbbd Mon Sep 17 00:00:00 2001
From: harshsoni2024 <64592571+harshsoni2024@users.noreply.github.com>
Date: Fri, 30 Aug 2024 09:49:07 +0530
Subject: [PATCH] Connector: vertexai schema (#17632)

* vertexai connection schema

* remove yaml file
---
 .../testConnections/mlmodel/vertexai.json     | 14 +++++++
 .../mlmodel/vertexaiConnection.json           | 38 +++++++++++++++++++
 .../entity/services/mlmodelService.json       |  8 +++-
 3 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 openmetadata-service/src/main/resources/json/data/testConnections/mlmodel/vertexai.json
 create mode 100644 openmetadata-spec/src/main/resources/json/schema/entity/services/connections/mlmodel/vertexaiConnection.json

diff --git a/openmetadata-service/src/main/resources/json/data/testConnections/mlmodel/vertexai.json b/openmetadata-service/src/main/resources/json/data/testConnections/mlmodel/vertexai.json
new file mode 100644
index 00000000000..f24b72fbbbe
--- /dev/null
+++ b/openmetadata-service/src/main/resources/json/data/testConnections/mlmodel/vertexai.json
@@ -0,0 +1,14 @@
+{
+    "name": "VertexAI",
+    "displayName": "VertexAI Test Connection",
+    "description": "Test Connection to validate the access against models and basic metadata extraction",
+    "steps": [
+      {
+        "name": "GetModels",
+        "description": "Validate that we can get the instances with the given credentials.",
+        "errorMessage": "Failed to get VertexAI models, please validate to the credentials of service account",
+        "shortCircuit": true,
+        "mandatory": true
+      }
+    ]
+}
\ No newline at end of file
diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/mlmodel/vertexaiConnection.json b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/mlmodel/vertexaiConnection.json
new file mode 100644
index 00000000000..992dea0fc81
--- /dev/null
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/mlmodel/vertexaiConnection.json
@@ -0,0 +1,38 @@
+{
+  "$id": "https://open-metadata.org/schema/entity/services/connections/mlmodel/vertexaiConnection.json",
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "VertexAIConnection",
+  "description": "Google VertexAI Connection Config",
+  "type": "object",
+  "javaType": "org.openmetadata.schema.services.connections.mlmodel.VertexAIConnection",
+  "definitions": {
+    "vertexAIType": {
+      "description": "Service type.",
+      "type": "string",
+      "enum": ["VertexAI"],
+      "default": "VertexAI"
+    }
+  },
+  "properties": {
+    "type": {
+      "title": "Service Type",
+      "description": "Service Type",
+      "$ref": "#/definitions/vertexAIType",
+      "default": "VertexAI"
+    },
+    "credentials": {
+      "title": "GCP Credentials",
+      "description": "GCP Credentials",
+      "$ref": "../../../../security/credentials/gcpCredentials.json"
+    },
+    "location": {
+      "title": "Project location",
+      "description": "location/region of google cloud project",
+      "type": "string"
+    }
+  },
+  "additionalProperties": false,
+  "required": [
+    "credentials", "location"
+  ]
+}
\ No newline at end of file
diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/services/mlmodelService.json b/openmetadata-spec/src/main/resources/json/schema/entity/services/mlmodelService.json
index 04c057e0677..4e114e5001c 100644
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/mlmodelService.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/mlmodelService.json
@@ -14,7 +14,7 @@
       "description": "Type of MlModel service",
       "type": "string",
       "javaInterfaces": ["org.openmetadata.schema.EnumInterface"],
-      "enum": ["Mlflow", "Sklearn", "CustomMlModel", "SageMaker"],
+      "enum": ["Mlflow", "Sklearn", "CustomMlModel", "SageMaker", "VertexAI"],
       "javaEnums": [
         {
           "name": "Mlflow"
@@ -27,6 +27,9 @@
         },
         {
           "name": "SageMaker"
+        },
+        {
+          "name": "VertexAI"
         }
       ]
     },
@@ -52,6 +55,9 @@
             },
             {
               "$ref": "./connections/mlmodel/sageMakerConnection.json"
+            },
+            {
+              "$ref": "./connections/mlmodel/vertexaiConnection.json"
             }
           ]
         }