Implement configurable default namespace for model references

cmd/mdltool/cli_test.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"os"
+	"os/exec"
+	"strings"
+	"testing"
+)
+
+func TestCLIDefaultRegistry(t *testing.T) {
+	// Build the mdltool binary for testing
+	buildCmd := exec.Command("go", "build", "-o", "test-mdltool", ".")
+	if err := buildCmd.Run(); err != nil {
+		t.Fatalf("Failed to build test binary: %v", err)
+	}
+	defer os.Remove("test-mdltool")
+
+	tests := []struct {
+		name     string
+		args     []string
+		contains string
+	}{
+		{
+			name:     "help shows default-registry option",
+			args:     []string{"--help"},
+			contains: "-default-registry",
+		},
+		{
+			name:     "version works",
+			args:     []string{"--version"},
+			contains: "model-distribution-tool version",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cmd := exec.Command("./test-mdltool", tt.args...)
+			output, err := cmd.CombinedOutput()
+
+			// For help and version, these should exit successfully
+			if tt.name == "help shows default-registry option" || tt.name == "version works" {
+				if err != nil {
+					t.Errorf("Command failed: %v, output: %s", err, output)
+				}
+			}
+
+			if !strings.Contains(string(output), tt.contains) {
+				t.Errorf("Expected output to contain %q, got: %s", tt.contains, output)
+			}
+		})
+	}
+}

cmd/mdltool/main.go

@@ -32,13 +32,15 @@ const (
 )
 
 var (
-	storePath string
-	showHelp  bool
-	showVer   bool
+	storePath       string
+	defaultRegistry string
+	showHelp        bool
+	showVer         bool
 )
 
 func init() {
 	flag.StringVar(&storePath, "store-path", defaultStorePath, "Path to the model store")
+	flag.StringVar(&defaultRegistry, "default-registry", "", "Default registry for model references (e.g., registry.example.com)")
 	flag.BoolVar(&showHelp, "help", false, "Show help")
 	flag.BoolVar(&showVer, "version", false, "Show version")
 }
@@ -69,6 +71,10 @@ func main() {
 		distribution.WithUserAgent("model-distribution-tool/" + version),
 	}
 
+	if defaultRegistry != "" {
+		clientOpts = append(clientOpts, distribution.WithDefaultRegistry(defaultRegistry))
+	}
+
 	if username := os.Getenv("DOCKER_USERNAME"); username != "" {
 		if password := os.Getenv("DOCKER_PASSWORD"); password != "" {
 			clientOpts = append(clientOpts, distribution.WithRegistryAuth(username, password))

distribution/client.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/sirupsen/logrus"
 
+	"github.com/docker/model-distribution/internal/naming"
 	"github.com/docker/model-distribution/internal/progress"
 	"github.com/docker/model-distribution/internal/store"
 	"github.com/docker/model-distribution/registry"
@@ -33,12 +34,13 @@ type Option func(*options)
 
 // options holds the configuration for a new Client
 type options struct {
-	storeRootPath string
-	logger        *logrus.Entry
-	transport     http.RoundTripper
-	userAgent     string
-	username      string
-	password      string
+	storeRootPath      string
+	logger             *logrus.Entry
+	transport          http.RoundTripper
+	userAgent          string
+	username           string
+	password           string
+	defaultRegistry    string
 }
 
 // WithStoreRootPath sets the store root path
@@ -87,6 +89,15 @@ func WithRegistryAuth(username, password string) Option {
 	}
 }
 
+// WithDefaultRegistry sets the default registry namespace for model references
+func WithDefaultRegistry(registry string) Option {
+	return func(o *options) {
+		if registry != "" {
+			o.defaultRegistry = registry
+		}
+	}
+}
+
 func defaultOptions() *options {
 	return &options{
 		logger:    logrus.NewEntry(logrus.StandardLogger()),
@@ -124,6 +135,13 @@ func NewClient(opts ...Option) (*Client, error) {
 		registryOpts = append(registryOpts, registry.WithAuthConfig(options.username, options.password))
 	}
 
+	// Add default registry namespace if provided
+	if options.defaultRegistry != "" {
+		registryOpts = append(registryOpts, registry.WithDefaultNamespace(options.defaultRegistry))
+		// Also set it globally for store operations
+		naming.SetDefaultNamespace(options.defaultRegistry)
+	}
+
 	options.logger.Infoln("Successfully initialized store")
 	return &Client{
 		store:    s,

distribution/integration_test.go

@@ -0,0 +1,76 @@
+package distribution
+
+import (
+	"testing"
+	"path/filepath"
+	"os"
+)
+
+func TestDefaultRegistryIntegration(t *testing.T) {
+	// Create a temporary directory for the test store
+	tempDir, err := os.MkdirTemp("", "test-store-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	tests := []struct {
+		name            string
+		defaultRegistry string
+		reference       string
+		expectError     bool
+		description     string
+	}{
+		{
+			name:            "no default registry - standard behavior",
+			defaultRegistry: "",
+			reference:       "library/alpine:latest",
+			expectError:     true, // will fail because registry is not accessible, but parsing should work
+			description:     "Should use Docker Hub as default",
+		},
+		{
+			name:            "custom default registry applied",
+			defaultRegistry: "registry.example.com",
+			reference:       "mymodel:latest",
+			expectError:     true, // will fail because registry is not accessible, but parsing should work
+			description:     "Should apply custom default registry",
+		},
+		{
+			name:            "explicit registry preserved",
+			defaultRegistry: "registry.example.com",
+			reference:       "other.registry.com/mymodel:latest",
+			expectError:     true, // will fail because registry is not accessible, but parsing should work  
+			description:     "Should preserve explicit registry even when default is set",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			storeDir := filepath.Join(tempDir, tt.name)
+
+			// Create client options
+			opts := []Option{
+				WithStoreRootPath(storeDir),
+			}
+			if tt.defaultRegistry != "" {
+				opts = append(opts, WithDefaultRegistry(tt.defaultRegistry))
+			}
+
+			// Create distribution client
+			client, err := NewClient(opts...)
+			if err != nil {
+				t.Fatalf("Failed to create client: %v", err)
+			}
+
+			// Test that the client can be created and the reference parsing works
+			// We expect these to fail with network errors since the registries don't exist,
+			// but we want to ensure the parsing is correct
+			_ = client
+
+			// For this test, we're mainly validating that the client can be created
+			// with the default registry configuration. The actual network operations
+			// would require a running registry which is beyond the scope of this test.
+			t.Logf("Successfully created client with default registry: %s", tt.defaultRegistry)
+		})
+	}
+}

internal/naming/README.md

@@ -0,0 +1,122 @@
+# Naming Package
+
+The naming package provides configurable default namespace support for container registry references in the model-distribution system.
+
+## Problem
+
+By default, the `github.com/google/go-containerregistry/pkg/name` package assumes Docker Hub (`index.docker.io`) as the default registry for image references that don't include an explicit registry hostname. For model distribution systems, it's often desirable to use a different default registry.
+
+## Solution
+
+This package provides wrapper functions around `name.ParseReference` and `name.ParseTag` that apply a configurable default namespace when no explicit registry is provided in the reference.
+
+## Usage
+
+### Client-specific Configuration
+
+```go
+import (
+    "github.com/docker/model-distribution/registry"
+    "github.com/docker/model-distribution/distribution"
+)
+
+// Configure registry client with custom default namespace
+registryClient := registry.NewClient(
+    registry.WithDefaultNamespace("registry.example.com"),
+)
+
+// Configure distribution client with custom default namespace
+distClient, err := distribution.NewClient(
+    distribution.WithStoreRootPath("/path/to/store"),
+    distribution.WithDefaultRegistry("registry.example.com"),
+)
+```
+
+### Global Configuration
+
+```go
+import "github.com/docker/model-distribution/internal/naming"
+
+// Set global default namespace
+naming.SetDefaultNamespace("registry.example.com")
+
+// Use global convenience functions
+ref, err := naming.ParseReference("mymodel:latest")
+// Will resolve to registry.example.com/mymodel:latest
+
+tag, err := naming.ParseTag("mymodel:latest") 
+// Will resolve to registry.example.com/mymodel:latest
+```
+
+### Direct Usage
+
+```go
+import "github.com/docker/model-distribution/internal/naming"
+
+// Create a namespace configuration
+ns := &naming.DefaultNamespace{Registry: "registry.example.com"}
+
+// Parse references with custom default
+ref, err := ns.ParseReference("mymodel:latest")
+// Will resolve to registry.example.com/mymodel:latest
+
+// Explicit registries are preserved
+ref, err := ns.ParseReference("other.registry.com/mymodel:latest")
+// Will remain other.registry.com/mymodel:latest
+```
+
+## CLI Usage
+
+The `mdltool` command-line interface supports a `--default-registry` flag:
+
+```bash
+# Use custom default registry
+mdltool --default-registry registry.example.com pull mymodel:latest
+
+# This will pull from registry.example.com/mymodel:latest instead of 
+# the Docker Hub default index.docker.io/library/mymodel:latest
+```
+
+## Behavior
+
+- **No explicit registry**: Default namespace is applied
+  - Input: `mymodel:latest` → Output: `registry.example.com/mymodel:latest`
+  - Input: `user/mymodel:latest` → Output: `registry.example.com/user/mymodel:latest`
+
+- **Explicit registry**: Default namespace is ignored
+  - Input: `other.registry.com/mymodel:latest` → Output: `other.registry.com/mymodel:latest`
+  - Input: `localhost:5000/mymodel:latest` → Output: `localhost:5000/mymodel:latest`
+
+- **No default configured**: Falls back to standard go-containerregistry behavior
+  - Input: `mymodel:latest` → Output: `index.docker.io/library/mymodel:latest`
+
+## Registry Detection
+
+The package uses heuristics to detect whether a reference already contains an explicit registry:
+
+- Contains a dot (`.`) before the first slash → Likely a registry hostname
+- Contains a colon (`:`) followed by digits before the first slash → Likely a registry with port
+- Contains a colon followed by non-digits → Likely a tag, not a registry port
+
+Examples:
+- `mymodel:latest` → No explicit registry (`:latest` is a tag)
+- `registry.com/mymodel:latest` → Has explicit registry (contains dot)
+- `localhost:5000/mymodel:latest` → Has explicit registry (port number)
+- `user/mymodel:latest` → No explicit registry
+
+## Integration Points
+
+This functionality is integrated at the following levels:
+
+1. **Registry Client**: `registry.Client` can be configured with a default namespace
+2. **Distribution Client**: `distribution.Client` configures both registry client and store operations
+3. **Store Operations**: Uses global namespace configuration for tag/reference parsing
+4. **CLI Tool**: Accepts `--default-registry` flag to configure the default
+
+## Backward Compatibility
+
+This change is fully backward compatible:
+
+- Existing code without default namespace configuration continues to work unchanged
+- All explicit registry references continue to work as before
+- The default behavior (Docker Hub) is preserved when no configuration is provided