From c01df6681e5f94d74dcc5fd85a9ea4af4b9dcea6 Mon Sep 17 00:00:00 2001
From: Yeesian Ng <ysian@google.com>
Date: Tue, 5 Mar 2024 09:02:20 -0800
Subject: [PATCH] feat: Allow module_name to be specified with the base logger.

PiperOrigin-RevId: 612860933
---
 google/cloud/aiplatform/base.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
index 385d407e63..cb4c8ecdf4 100644
--- a/google/cloud/aiplatform/base.py
+++ b/google/cloud/aiplatform/base.py
@@ -92,6 +92,8 @@ def log_create_complete(
         cls: Type["VertexAiResourceNoun"],
         resource: proto.Message,
         variable_name: str,
+        *,
+        module_name: str = "aiplatform",
     ):
         """Logs create event is complete.
 
@@ -102,17 +104,23 @@ def log_create_complete(
                 Vertex AI Resource Noun class that is being created.
             resource (proto.Message):
                 Vertex AI Resource proto.Message
-            variable_name (str): Name of variable to use for code snippet
+            variable_name (str):
+                Name of variable to use for code snippet.
+            module_name (str):
+                The module namespace under which the Vertex AI Resource Noun
+                is available. Defaults to `aiplatform`.
         """
         self.info(f"{cls.__name__} created. Resource name: {resource.name}")
         self.info(f"To use this {cls.__name__} in another session:")
-        self.info(f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')")
+        self.info(f"{variable_name} = {module_name}.{cls.__name__}('{resource.name}')")
 
     def log_create_complete_with_getter(
         self,
         cls: Type["VertexAiResourceNoun"],
         resource: proto.Message,
         variable_name: str,
+        *,
+        module_name: str = "aiplatform",
     ):
         """Logs create event is complete.
 
@@ -123,11 +131,16 @@ def log_create_complete_with_getter(
                 Vertex AI Resource Noun class that is being created.
             resource (proto.Message):
                 Vertex AI Resource proto.Message
-            variable_name (str): Name of variable to use for code snippet
+            variable_name (str):
+                Name of variable to use for code snippet.
+            module_name (str):
+                The module namespace under which the Vertex AI Resource Noun
+                is available. Defaults to `aiplatform`.
         """
         self.info(f"{cls.__name__} created. Resource name: {resource.name}")
         self.info(f"To use this {cls.__name__} in another session:")
-        self.info(f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')")
+        usage_message = f"{module_name}.{cls.__name__}.get('{resource.name}')"
+        self.info(f"{variable_name} = {usage_message}")
 
     def log_delete_with_lro(
         self,