Add logic on nbc-webhook to pickup internal or external registry

components/odh-notebook-controller/controllers/notebook_webhook.go

@@ -20,6 +20,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"sort"
+	"strings"
 
 	"github.com/go-logr/logr"
 	nbv1 "github.com/kubeflow/kubeflow/components/notebook-controller/api/v1"
@@ -28,7 +30,10 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/client-go/dynamic"
+	"k8s.io/client-go/rest"
 	"k8s.io/utils/pointer"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
@@ -40,6 +45,7 @@ import (
 type NotebookWebhook struct {
 	Log         logr.Logger
 	Client      client.Client
+	Config      *rest.Config
 	Decoder     *admission.Decoder
 	OAuthConfig OAuthConfig
 }
@@ -246,6 +252,12 @@ func (w *NotebookWebhook) Handle(ctx context.Context, req admission.Request) adm
 		if err != nil {
 			return admission.Errored(http.StatusInternalServerError, err)
 		}
+
+		// Check Imagestream Info
+		err = SetContainerImageFromRegistry(ctx, w.Config, notebook, log)
+		if err != nil {
+			return admission.Errored(http.StatusInternalServerError, err)
+		}
 	}
 
 	// Inject the OAuth proxy if the annotation is present but only if Service Mesh is disabled
@@ -438,3 +450,106 @@ func InjectCertConfig(notebook *nbv1.Notebook, configMapName string) error {
 	}
 	return nil
 }
+
+// SetContainerImageFromRegistry checks if there is an internal registry and takes the corresponding actions to set the container.image value.
+// If an internal registry is detected, it uses the default values specified in the Notebook Custom Resource (CR).
+// Otherwise, it checks the last-image-selection annotation to find the image stream and fetches the image from status.dockerImageReference,
+// assigning it to the container.image value.
+func SetContainerImageFromRegistry(ctx context.Context, config *rest.Config, notebook *nbv1.Notebook, log logr.Logger) error {
+    // Create a dynamic client
+    dynamicClient, err := dynamic.NewForConfig(config)
+    if err != nil {
+        log.Error(err, "Error creating dynamic client")
+        return err
+    }
+    // Specify the GroupVersionResource for imagestreams
+    ims := schema.GroupVersionResource{
+        Group:    "image.openshift.io",
+        Version:  "v1",
+        Resource: "imagestreams",
+    }
+
+    annotations := notebook.GetAnnotations()
+    if annotations != nil {
+        if imageSelection, exists := annotations["notebooks.opendatahub.io/last-image-selection"]; exists {
+            // Check if the image selection has an internal registry, if so  will pickup this. This value constructed on the initialization of the Notebook CR.
+            if strings.Contains(notebook.Spec.Template.Spec.Containers[0].Image, "image-registry.openshift-image-registry.svc:5000") {
+                log.Info("Internal registry found. Will pickup the default value from image field.")
+                return nil
+            } else {
+                // Split the imageSelection to imagestream and tag
+                parts := strings.Split(imageSelection, ":")
+                if len(parts) != 2 {
+                    log.Error(nil, "Invalid image selection format")
+                    return fmt.Errorf("invalid image selection format")
+                }
+
+                imagestreamName := parts[0]
+                tag := parts[1]
+
+                // Specify the namespaces to search in
+                namespaces := []string{"opendatahub", "redhat-ods-applications"}
+
+                imagestreamFound := false
+
+                for _, namespace := range namespaces {
+                    // List imagestreams in the specified namespace
+                    imagestreams, err := dynamicClient.Resource(ims).Namespace(namespace).List(ctx, metav1.ListOptions{})
+                    if err != nil {
+                        log.Error(err, "Cannot list imagestreams", "namespace", namespace)
+                        continue
+                    }
+
+                    // Iterate through the imagestreams to find matches
+                    for _, item := range imagestreams.Items {
+                        metadata := item.Object["metadata"].(map[string]interface{})
+                        name := metadata["name"].(string)
+
+                        if name == imagestreamName {
+                            status := item.Object["status"].(map[string]interface{})
+
+                            log.Info("No Internal registry found, pick up imageHash from status.tag.dockerImageReference")
+
+                            tags := status["tags"].([]interface{})
+                            for _, t := range tags {
+                                tagMap := t.(map[string]interface{})
+                                tagName := tagMap["tag"].(string)
+                                if tagName == tag {
+                                    items := tagMap["items"].([]interface{})
+                                    if len(items) > 0 {
+                                        // Sort items by creationTimestamp to get the most recent one
+                                        sort.Slice(items, func(i, j int) bool {
+                                            iTime := items[i].(map[string]interface{})["created"].(string)
+                                            jTime := items[j].(map[string]interface{})["created"].(string)
+                                            return iTime > jTime // Lexicographical comparison of RFC3339 timestamps
+                                        })
+                                        imageHash := items[0].(map[string]interface{})["dockerImageReference"].(string)
+                                        notebook.Spec.Template.Spec.Containers[0].Image = imageHash
+                                        // Update the JUPYTER_IMAGE environment variable
+                                        for i, envVar := range notebook.Spec.Template.Spec.Containers[0].Env {
+                                            if envVar.Name == "JUPYTER_IMAGE" {
+                                                notebook.Spec.Template.Spec.Containers[0].Env[i].Value = imageHash
+                                                break
+                                            }
+                                        }
+                                        imagestreamFound = true
+                                        break
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    if imagestreamFound {
+                        break
+                    }
+                }
+
+                if !imagestreamFound {
+                    log.Info("Imagestream not found in any of the specified namespaces", "imagestreamName", imagestreamName, "tag", tag)
+                }
+            }
+        }
+    }
+
+    return nil
+}

components/odh-notebook-controller/controllers/suite_test.go

@@ -1,5 +1,4 @@
 /*
-
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
@@ -141,6 +140,7 @@ var _ = BeforeSuite(func() {
 		Handler: &NotebookWebhook{
 			Log:    ctrl.Log.WithName("controllers").WithName("notebook-controller"),
 			Client: mgr.GetClient(),
+			Config: mgr.GetConfig(),
 			OAuthConfig: OAuthConfig{
 				ProxyImage: OAuthProxyImage,
 			},

components/odh-notebook-controller/main.go

@@ -116,6 +116,7 @@ func main() {
 		Handler: &controllers.NotebookWebhook{
 			Log:    ctrl.Log.WithName("controllers").WithName("Notebook"),
 			Client: mgr.GetClient(),
+			Config: mgr.GetConfig(),
 			OAuthConfig: controllers.OAuthConfig{
 				ProxyImage: oauthProxyImage,
 			},