Ivo Bardarov

"""
DAG to refresh Docker ECR authentication token
Updates the docker_default connection with fresh ECR credentials using JWT authentication

You should not have your own: ~/.docker/config.json

"""

import base64
import logging
from datetime import datetime, timedelta
from typing import Any, Dict

import boto3
import requests
from airflow.decorators import dag, task
from airflow.hooks.base import BaseHook
from airflow.models.variable import Variable
from airflow.providers.docker.hooks.docker import DockerHook

logger = logging.getLogger("ecr_docker_token_refresh")

ecr_aws_account = Variable.get("ecr_aws_account")
ecr_aws_region_name = Variable.get("ecr_aws_region_name")

default_args = {
    "retry_delay": timedelta(minutes=1),
    "depends_on_past": False,
    "retries": 2,
    "email_on_failure": False,
    "email_on_retry": False,
    "queue": "systemqueue",
    "pool": "systemqueue",
}

connection_id = "docker_default"
airflow_api_connection_id = "airflow-api"


def get_jwt_token(endpoint_url: str, username: str, password: str) -> str:
    """Get JWT token from Airflow API"""
    auth_url = f"{endpoint_url}/auth/token"
    payload = {"username": username, "password": password}
    headers = {"Content-Type": "application/json"}
    logger.info(f"Requesting JWT token from {auth_url}")

    response = requests.post(auth_url, json=payload, headers=headers)
    response.raise_for_status()

    token_data = response.json()
    access_token = token_data.get("access_token")

    if not access_token:
        raise ValueError("No access_token found in response")

    logger.info("Successfully obtained JWT token")
    return access_token


def update_connection_password_with_jwt(
    endpoint_url: str, jwt_token: str, password: str
) -> bool:
    """Update connection password using JWT token with v2 bulk API"""
    url = f"{endpoint_url}/api/v2/connections"

    # First, get the current connection to preserve other fields
    get_url = f"{endpoint_url}/api/v2/connections/{connection_id}"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {jwt_token}",
    }

    logger.info(f"Getting current connection {connection_id}")

    try:
        get_response = requests.get(get_url, headers=headers)
        get_response.raise_for_status()
        current_connection = get_response.json()

        logger.info(f"Current connection retrieved successfully")

        # Prepare bulk update payload using v2 API
        payload = {
            "actions": [
                {
                    "action": "update",
                    "entities": [
                        {
                            "connection_id": connection_id,
                            "conn_type": current_connection.get("conn_type", "docker"),
                            "password": password,  # This is what we're updating
                        }
                    ],
                    "action_on_non_existence": "fail",
                }
            ]
        }

        logger.info(f"Updating connection {connection_id} at {url}")

        response = requests.patch(url, json=payload, headers=headers)
        response.raise_for_status()

        response_data = response.json()
        logger.info(f"Bulk update response: {response_data}")

        # Check if update was successful
        update_results = response_data.get("update", {})
        success_count = len(update_results.get("success", []))
        error_count = len(update_results.get("errors", []))

        if success_count > 0 and error_count == 0:
            logger.info("Connection password updated successfully")
            return True
        else:
            logger.error(
                f"Update failed - Success: {success_count}, Errors: {error_count}"
            )
            if error_count > 0:
                logger.error(f"Errors: {update_results.get('errors', [])}")
            return False

    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to update connection: {e}")
        if hasattr(e, "response") and e.response is not None:
            logger.error(f"Response status: {e.response.status_code}")
            logger.error(f"Response text: {e.response.text}")
        raise


@dag(
    default_args=default_args,
    schedule="55 5,17 * * *",
    start_date=datetime.now() - timedelta(days=1),
    max_active_runs=1,
    catchup=False,
    tags=["airflow", "docker", "ecr"],
    dag_id="refresh_docker_token_v4",
    description="Refresh Docker ECR token using JWT authentication",
)
def refresh_docker_token():
    @task(priority_weight=5, pool="systemqueue")
    def extract_ecr_token() -> Dict[str, Any]:
        """Extract ECR authorization token using boto3"""
        logger.info("Starting ECR token extraction")

        try:
            logger.info(f"Connecting to ECR in region {ecr_aws_region_name}")
            ecr_client = boto3.client("ecr", region_name=ecr_aws_region_name)

            logger.info(f"Requesting authorization token for account {ecr_aws_account}")
            response = ecr_client.get_authorization_token(registryIds=[ecr_aws_account])

            auth_data = response["authorizationData"][0]
            token = auth_data["authorizationToken"]
            registry_url = auth_data["proxyEndpoint"]
            expires_at = auth_data["expiresAt"]

            logger.info("Successfully retrieved token")
            logger.info(f"Registry URL: {registry_url}")
            logger.info(f"Token expires at: {expires_at}")

            decoded_token = base64.b64decode(token).decode()
            username, password = decoded_token.split(":", 1)

            logger.info(f"Decoded username: {username}")

            return {
                "registry_url": registry_url,
                "username": username,
                "password": password,
                "expires_at": expires_at.isoformat(),
                "raw_token": token,
            }

        except Exception as e:
            logger.error(f"Failed to extract ECR token: {str(e)}")
            raise

    @task(priority_weight=5, pool="systemqueue")
    def update_docker_connection(token_data: Dict[str, Any]) -> str:
        """Update Docker connection using JWT authentication"""
        logger.info("Starting Docker connection update using JWT authentication")
        logger.info("Token data received from previous task")

        try:
            # Get Airflow API connection details
            logger.info(
                f"Retrieving Airflow API connection: {airflow_api_connection_id}"
            )
            api_connection = BaseHook.get_connection(airflow_api_connection_id)

            endpoint_url = f"{api_connection.schema}://{api_connection.host}"
            if api_connection.port:
                endpoint_url += f":{api_connection.port}"

            username = api_connection.login
            password = api_connection.password

            logger.info(f"Using endpoint: {username} @ {endpoint_url}")
            jwt_token = get_jwt_token(endpoint_url, username, password)

            success = update_connection_password_with_jwt(
                endpoint_url, jwt_token, token_data["password"]
            )

            if success:
                return "SUCCESS: Docker connection updated successfully using JWT authentication"
            else:
                raise Exception("Failed to update connection")

        except Exception as e:
            logger.error(f"Failed to update Docker connection: {str(e)}")
            raise

    @task(priority_weight=3, pool="systemqueue")
    def test_docker_connection() -> str:
        """Test the updated Docker connection"""
        logger.info("Testing DockerHook...")

        try:
            # First get the connection details to debug
            connection = BaseHook.get_connection("docker_default")

            docker_hook = DockerHook(
                docker_conn_id="docker_default", base_url="unix://var/run/docker.sock"
            )
            logger.info("DockerHook created successfully")

            # Try to get docker client (this will test the connection more thoroughly)
            docker_client = docker_hook.get_conn()
            logger.info("Docker client connection established", docker_client.version())
            return "SUCCESS: Docker connection tested and working with DockerHook"

        except Exception as client_error:
            logger.error(f"Docker client test failed: {client_error}")

            # Show connection properties on failure
            try:
                connection = BaseHook.get_connection("docker_default")
 
            except Exception as conn_error:
                logger.error(f"Could not retrieve connection properties: {conn_error}")

            return f"FAILED: Docker connection test failed: {client_error}"

    # Task flow
    token_data = extract_ecr_token()
    update_result = update_docker_connection(token_data)
    test_result = test_docker_connection()

    token_data >> update_result >> test_result


refresh_docker_token_dag = refresh_docker_token()

sudo apt install -y postgresql-common

sudo /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh

apt install postgresql

pg_lsclusters
dpkg --get-selections | grep postgres

pg_renamecluster 17 main main_pristine

sudo pg_upgradecluster 13 main

# pg_lsclusters
Ver Cluster       Port Status Owner    Data directory                       Log file
13  main          5432 online postgres /var/lib/postgresql/13/main          /var/log/postgresql/postgresql-13-main.log
17  main_pristine 5434 online postgres /var/lib/postgresql/17/main_pristine /var/log/postgresql/postgresql-17-main_pristine.log

# sudo pg_upgradecluster 13 main
.... a lot of upgrading

# pg_lsclusters
Ver Cluster       Port Status Owner    Data directory                       Log file
13  main          5435 down   postgres /var/lib/postgresql/13/main          /var/log/postgresql/postgresql-13-main.log
17  main          5432 online postgres /var/lib/postgresql/17/main          /var/log/postgresql/postgresql-17-main.log
17  main_pristine 5434 online postgres /var/lib/postgresql/17/main_pristine /var/log/postgresql/postgresql-17-main_pristine.log

pg_dropcluster 13 main --stop
pg_dropcluster 17 main_pristine --stop

  ls -l /home/user/projects/pizza/seed/images/themes/pizza/
  ls -l /home/user/projects/pizza/seed/images/themes/
  ls -l /home/user/projects/pizza/seed/images/
  ls -l /home/user/projects/pizza/seed/
  ls -l /home/user/projects/pizza/

ls -l/home/user/projects/pizza/seeds/images/themes/pizza/01.jpg

meld project-name/ project-name-other/

cdf /home/user/some-project/some-folder/file_name.extension

cdf() {
    if [ $# -eq 0 ]; then
        echo "No file path provided."
        return 1
    fi

    # Join all arguments with spaces
    local full_path="$*"


    if [ -f "$full_path" ]; then
        # If it's a file, extract the directory path
        local dir_path
        dir_path=$(dirname "$full_path")
    elif [ -d "$full_path" ]; then
        # If it's a directory, use it directly
        local dir_path="$full_path"
    else
        echo "The path provided is neither a file nor a directory."
        return 1
    fi


    # Change to the directory
    cd "$dir_path" || {
        echo "Failed to change directory to $dir_path"
        return 1
    }
}

# Snowflake account
export SNOWFLAKE_ACCOUNT=YOURACCOUNT.us-east-1
export SNOWFLAKE_WAREHOUSE=YOURWAREHOUSE

export DBT_PROFILE=development
export DBT_ROLE=DEVELOPMENT

# Tenant / Used for the query tagging
export TENANT_NAME=CLIENT_NAME

# Source
export DBT_SOURCE_DATABASE=RAW_ZONE
export DBT_SOURCE_SCHEMA=${TENANT_NAME}

# Target
export DBT_TARGET_DATABASE=STANDARD_ZONE_DEVELOPMENT
export DBT_TARGET_SCHEMA=${TENANT_NAME}_${SNOWFLAKE_USER}

development:
  target: development
  outputs:
    development:
      type: snowflake
      account: "{{ env_var('SNOWFLAKE_ACCOUNT') }}"
      warehouse: "{{ env_var('SNOWFLAKE_WAREHOUSE') }}"
      user: "{{ env_var('SNOWFLAKE_USER') }}"
      password: "{{ env_var('SNOWFLAKE_PASSWORD') }}"
      role: "{{env_var('DBT_ROLE')}}"
      database: "{{ env_var('DBT_TARGET_DATABASE') }}"
      schema: "{{ env_var('DBT_TARGET_SCHEMA') }}"
      threads: 8


production:
  target: production
  outputs:
    production:
      type: snowflake
      account: "{{ env_var('SNOWFLAKE_ACCOUNT') }}"
      user: "{{ env_var('SNOWFLAKE_USER') }}"
      password: "{{ env_var('SNOWFLAKE_PASSWORD') }}"
      role: "{{env_var('DBT_ROLE')}}"
      warehouse: "{{ env_var('SNOWFLAKE_WAREHOUSE') }}"
      database: "{{ env_var('DBT_TARGET_DATABASE') }}"
      schema: "{{ env_var('DBT_TARGET_SCHEMA') }}"
      threads: 8

profile: "{{ env_var('DBT_PROFILE', 'development') }}"

vars:
  source_database: "{{ env_var('DBT_SOURCE_DATABASE') }}"
  source_schema: "{{ env_var('DBT_SOURCE_SCHEMA') }}"

  export DBT_PROFILES_DIR="."
  dbt --no-anonymous-usage-stats run

dbt:
  #!/bin/bash
  echo "SOURCE: $DBT_SOURCE_DATABASE / $DBT_SOURCE_SCHEMA"
  echo "TARGET: $DBT_TARGET_DATABASE / $DBT_TARGET_SCHEMA"
  DBT_PROFILES_DIR="{{invocation_directory()}}"

  poetry run \
    dbt --no-anonymous-usage-stats run \
    --fail-fast

"dbt.profilesDirOverride": "${workspaceFolder}",

  "files.associations": {
        "*.sql": "jinja-sql"
    },