# Copyright 2022 Google LLC
#
# 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
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

---
apiVersion: v1
kind: ServiceAccount
metadata:
  namespace: "default"
  name: pgpool-operator
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  name: pgpool-operator-cluster-role
rules:
  - apiGroups: [""]
    resources:
    - namespaces
    verbs: ["*"]
  - apiGroups: [apiextensions.k8s.io]
    resources: [customresourcedefinitions]
    verbs: [watch, list]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
  namespace: "default"
  name: pgpool-operator-role
rules:
  - apiGroups: ["apps"]
    resources: ["deployments", "statefulsets"]
    verbs: ["*"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: pgpool-operator-cluster-rolebinding
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: pgpool-operator-cluster-role
subjects:
  - kind: ServiceAccount
    name: pgpool-operator
    namespace: "default"
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
  namespace: "default"
  name: pgpool-operator-rolebinding
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: Role
  name: pgpool-operator-role
subjects:
  - kind: ServiceAccount
    namespace: "default"
    name: pgpool-operator
---
apiVersion: v1
kind: ConfigMap
metadata:
  name: pgpool-operator-script
data:
  pgpool.py: |
    # Copyright 2022 Google LLC
    #
    # 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
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.

    import asyncio
    import datetime
    import logging
    import kopf
    from kubernetes import client
    from kubernetes.client.rest import ApiException

    LOCK: asyncio.Lock

    @kopf.on.startup()
    async def startup(**_):
        """
        uses the running asyncio loop by default
        """
        global LOCK
        LOCK = asyncio.Lock()


    @kopf.on.startup()
    def configure(settings: kopf.OperatorSettings, **_):
        settings.posting.level = logging.WARNING
        settings.watching.connect_timeout = 1 * 60
        settings.watching.server_timeout = 10 * 60


    @kopf.on.probe(id='now')
    def get_current_timestamp(**kwargs):
        return datetime.datetime.utcnow().isoformat()


    @kopf.on.login()
    def login(**kwargs):
        global api

        conn = kopf.login_via_client(**kwargs)
        api = client.AppsV1Api()
        return conn


    def replicas_changed(old, new, **_):
        new_replicas = new.get('spec', {}).get('replicas', 0) if new else 0
        old_replicas = old.get('spec', {}).get('replicas', 0) if old else 0
        return new_replicas != old_replicas



    @kopf.on.update(kind="StatefulSet",
                    when=replicas_changed,
                    labels={
                        "app.kubernetes.io/component": "postgresql",
                        "app.kubernetes.io/instance": "accounts-db",
                    })
    def reconcile_backend_nodes(logger, namespace, new, **_):
        replicas = new.get('spec', {}).get('replicas', 0) if new else 0
        hosts = [
            f"{i}:accounts-db-postgresql-{i}.accounts-db-postgresql-headless:5432" \
            for i in range(replicas)
        ]

        def propagate_hostenv(envvar, hosts):
            if envvar.name == "PGPOOL_BACKEND_NODES":
                return {
                    "name": "PGPOOL_BACKEND_NODES",
                    "value": ",".join(hosts),
                }
            return envvar

        try:
            pgpool = api.read_namespaced_deployment(name="accounts-db-pgpool", namespace=namespace)
            for container in pgpool.spec.template.spec.containers:
                container.env = [propagate_hostenv(envvar, hosts) for envvar in container.env]

            api.patch_namespaced_deployment(
                name="accounts-db-pgpool",
                namespace=namespace,
                body=pgpool
            )
            logger.info("PGPool deployment updated")
        except ApiException as e:
            raise kopf.TemporaryError("Error when calling AppsV1Api: %s\n" % e, delay=60)

---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: pgpool-operator
spec:
  replicas: 1
  strategy:
    type: Recreate
  selector:
    matchLabels:
      app: pgpool-operator
  template:
    metadata:
      labels:
        app: pgpool-operator
    spec:
      serviceAccountName: pgpool-operator
      containers:
        - name: operator
          image: python:3.12-bullseye@sha256:deda19eaa3781e480eeb65f0463cf8d84874c4aa8c8997cdcb4d12c4ff2db957
          env:
          - name: NAMESPACE
            valueFrom:
              fieldRef:
                fieldPath: metadata.namespace
          command: ["/bin/sh", "-c", "pip3 install kopf kubernetes && kopf run /app/pgpool.py --standalone --namespace=$(NAMESPACE) --log-format=plain --liveness=http://0.0.0.0:8080/healthz"]
          volumeMounts:
            - name: pgpool-operator-script
              mountPath: /app
          livenessProbe:
            httpGet:
              path: /healthz
              port: 8080
            initialDelaySeconds: 60
            periodSeconds: 30
            timeoutSeconds: 30
          readinessProbe:
            httpGet:
              path: /healthz
              port: 8080
            initialDelaySeconds: 60
            periodSeconds: 30
            timeoutSeconds: 30
          resources:
            requests:
              cpu: 250m
              memory: 512Mi
              ephemeral-storage: 1Gi
            limits:
              cpu: 250m
              memory: 512Mi
              ephemeral-storage: 1Gi
      volumes:
        - name: pgpool-operator-script
          configMap:
            name: pgpool-operator-script
