Merge pull request #12 from mikkeloscar/ownerreferences

Add owner references to PDBs

controller.go

@@ -11,7 +11,6 @@ import (
 	"k8s.io/client-go/pkg/api/v1"
 	"k8s.io/client-go/pkg/apis/apps/v1beta1"
 	pv1beta1 "k8s.io/client-go/pkg/apis/policy/v1beta1"
-	"k8s.io/client-go/rest"
 )
 
 const (
@@ -32,17 +31,7 @@ type PDBController struct {
 }
 
 // NewPDBController initializes a new PDBController.
-func NewPDBController(interval time.Duration, pdbNameSuffix string) (*PDBController, error) {
-	config, err := rest.InClusterConfig()
-	if err != nil {
-		return nil, err
-	}
-
-	client, err := kubernetes.NewForConfig(config)
-	if err != nil {
-		return nil, err
-	}
-
+func NewPDBController(interval time.Duration, client kubernetes.Interface, pdbNameSuffix string) (*PDBController, error) {
 	controller := &PDBController{
 		Interface:     client,
 		interval:      interval,
@@ -176,6 +165,14 @@ func (n *PDBController) addPDBs(namespace *v1.Namespace) error {
 			labels[heritageLabel] = pdbController
 			pdb.Name = r.Name
 			pdb.Namespace = r.Namespace
+			pdb.OwnerReferences = []metav1.OwnerReference{
+				{
+					APIVersion: "apps/v1",
+					Kind:       "Deployment",
+					Name:       r.Name,
+					UID:        r.UID,
+				},
+			}
 			pdb.Labels = labels
 			pdb.Spec.Selector = r.Spec.Selector
 		case v1beta1.StatefulSet:
@@ -186,6 +183,14 @@ func (n *PDBController) addPDBs(namespace *v1.Namespace) error {
 			labels[heritageLabel] = pdbController
 			pdb.Name = r.Name
 			pdb.Namespace = r.Namespace
+			pdb.OwnerReferences = []metav1.OwnerReference{
+				{
+					APIVersion: "apps/v1",
+					Kind:       "StatefulSet",
+					Name:       r.Name,
+					UID:        r.UID,
+				},
+			}
 			pdb.Labels = labels
 			pdb.Spec.Selector = r.Spec.Selector
 		}

main.go

@@ -1,12 +1,15 @@
 package main
 
 import (
+	"net/url"
 	"os"
 	"os/signal"
 	"syscall"
 	"time"
 
 	"gopkg.in/alecthomas/kingpin.v2"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 
 	log "github.com/sirupsen/logrus"
 )
@@ -18,13 +21,15 @@ const (
 
 type config struct {
 	Interval      time.Duration
+	APIServer     *url.URL
 	Debug         bool
 	PDBNameSuffix string
 }
 
 func main() {
 	config := config{}
 	kingpin.Flag("interval", "Interval between creating PDBs.").Default(defaultInterval).DurationVar(&config.Interval)
+	kingpin.Flag("apiserver", "API server url.").URLVar(&config.APIServer)
 	kingpin.Flag("debug", "Enable debug logging.").BoolVar(&config.Debug)
 	kingpin.Flag("pdb-name-suffix", "Specify default PDB name suffix.").Default(defaultPDBNameSuffix).StringVar(&config.PDBNameSuffix)
 	kingpin.Parse()
@@ -33,7 +38,26 @@ func main() {
 		log.SetLevel(log.DebugLevel)
 	}
 
-	controller, err := NewPDBController(config.Interval, config.PDBNameSuffix)
+	var err error
+	var kubeConfig *rest.Config
+
+	if config.APIServer != nil {
+		kubeConfig = &rest.Config{
+			Host: config.APIServer.String(),
+		}
+	} else {
+		kubeConfig, err = rest.InClusterConfig()
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+
+	client, err := kubernetes.NewForConfig(kubeConfig)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	controller, err := NewPDBController(config.Interval, client, config.PDBNameSuffix)
 	if err != nil {
 		log.Fatal(err)
 	}