scripts/main.py - master - mitinrs/jira

scripts/main.py
342 строки16 KB
Roman Mitin
flask
17 мар 2024, 15:34
17 мар 2024, 15:344517f35
import json  # Импорт модуля для работы с JSON данными
import psycopg2  # Импорт модуля для работы с PostgreSQL базой данных
import argparse  # Импорт модуля для парсинга аргументов командной строки
from psycopg2.extras import Json  # Импорт Json-адаптера для psycopg2, который позволяет работать с JSON данными в PostgreSQL
import psycopg2
from datetime import datetime, timezone
from jira import JIRA  # Импорт модуля для работы с API Jira
from configparser import ConfigParser  # Импорт модуля для работы с конфигурационными файлами
from datetime import datetime, timedelta  # Импорт модуля для работы с датами и временем
import requests  # Импорт модуля для выполнения HTTP-запросов
from requests.auth import HTTPBasicAuth  # Импорт класса для аутентификации HTTPBasicAuth
from tqdm import tqdm  # Импорт модуля для создания индикатора выполнения прогресса
from dateutil.relativedelta import relativedelta  # Импорт класса для работы с относительной датой и временем
import os  # Импорт модуля для работы с операционной системой
import socket
from db_to_parquet import export2parquet

# Счетчики и словарь для статистики по проектам
new_issues_count = 0
updated_issues_count = 0
projects_statistics = {}  # Словарь для хранения статистики в формате {название_проекта: {new: количество_новых, updated: количество_обновленных}}

# Initialize the ConfigParser and read the config file
config = ConfigParser()
config.read("config.ini")

# Connection parameters for JIRA and PostgreSQL from config.ini
JIRA_URL = config.get("JIRA", "URL").strip("'")
JIRA_EMAIL = config.get("JIRA", "EMAIL").strip("'")
JIRA_API_TOKEN = config.get("JIRA", "API_TOKEN").strip("'")
JIRA_PROJECTS = config.get("JIRA", "PROJECTS").strip().split(",")

DBNAME = config.get("DATABASE", "DATABASE").strip("'")
DBUSER = config.get("DATABASE", "USER").strip("'")
DBPASSWORD = config.get("DATABASE", "PASSWORD").strip("'")
DBHOST = config.get("DATABASE", "HOST").strip("'")

def log_script_execution(conn, log_details):
    with conn.cursor() as cursor:
        query = """
            INSERT INTO script_logs (
                id_comp, jql, duration, log_level, log_message, script_params, 
                start_time, end_time, total_issues_fetched, total_issues_new, 
                total_issues_updated, errors
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        cursor.execute(query, (
            log_details.get('id_comp'),
            log_details.get('jql'),
            log_details.get('duration'),
            log_details.get('log_level', 'INFO'),  # Default to INFO if not specified
            log_details.get('log_message'),
            log_details.get('script_params'),
            log_details.get('start_time'),
            log_details.get('end_time'),
            log_details.get('total_issues_fetched', 0),
            log_details.get('total_issues_new', 0),
            log_details.get('total_issues_updated', 0),
            log_details.get('errors', '')
        ))
    conn.commit()

#Определяем дату начала диапазона запроса
def get_relative_date(date, symbol):
    """
    Возвращает дату, измененную на заданный период назад от указанной даты.
    :param date: datetime - Исходная дата
    :param symbol: str - Символ периода ('d' - день, 'w' - неделя, 'm' - месяц, 'q' - квартал, 'y' - год, '' - без изменений)
    :return: datetime - Измененная дата
    """
    if symbol == 'd':
        return date - timedelta(days=1)
    elif symbol == 'w':
        return date - timedelta(weeks=1)
    elif symbol == 'm':
        return date - relativedelta(months=1)
    elif symbol == 'q':
        return date - relativedelta(months=3)
    elif symbol == 'y':
        return date - relativedelta(years=1)
    elif symbol == 'n':
        return date - relativedelta(years=5)
    elif symbol == '':
        return date
    else:
        raise ValueError("Неизвестный символ периода")

#Построение JQL запроса
def build_jql_query(projects, start_date=None, end_date=None):
    """
    Строит JQL запрос на основе списка проектов и опциональных дат начала и завершения.
    
    :param projects: Список проектов JIRA
    :param start_date: Необязательная начальная дата фильтрации (строка в формате YYYY-MM-DD)
    :param end_date: Необязательная конечная дата фильтрации (строка в формате YYYY-MM-DD)
    :return: Строка с JQL запросом
    """
    # Формируем базовую часть запроса
    jql_query ='project in (' + ', '.join(['"{}"'.format(proj) for proj in JIRA_PROJECTS]) + ')'
    
    # Добавляем фильтрацию по дате, если даты предоставлены
    date_filter_parts = []
    if start_date:
        date_filter_parts.append(f"updated >= '{start_date}'")
    #Доработать.Неверно фильтрует 
    #if end_date:
    #    date_filter_parts.append(f"updated <= '{end_date}'")

    # Если есть условия фильтрации по дате, добавляем их к запросу
    if date_filter_parts:
        jql_query += " AND " + " AND ".join(date_filter_parts)
    
    # Заканчиваем запрос сортировкой
    jql_query += " order by RANK"

    return jql_query

def upsert_task(issue, cursor):
    jira_key = issue["key"]
    raw_data = issue
    updated_at_jira = datetime.strptime(issue["fields"]["updated"], "%Y-%m-%dT%H:%M:%S.%f%z")

    cursor.execute("SELECT task_id, updated_at FROM jira_tasks WHERE jira_key = %s", (jira_key,))
    result = cursor.fetchone()

    if result:
        task_id, updated_at_db = result
        if updated_at_jira > updated_at_db:
            update_task(task_id, issue, updated_at_jira, raw_data, cursor)
    else:
        insert_new_task(issue, updated_at_jira, raw_data, cursor)

def insert_new_task(issue, updated_at, raw_data, cursor):
    global new_issues_count, projects_statistics
    assignee_name = issue["fields"].get("assignee", {}).get("displayName", "") if issue["fields"].get("assignee") else ""
    project_name = issue["fields"]["project"]["name"]
    cursor.execute("""
        INSERT INTO jira_tasks (jira_key, summary, description, status, priority, project, assignee, updated_at, raw_data)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING task_id
    """, (
        issue["key"], issue["fields"]["summary"], issue["fields"].get("description"),
        issue["fields"]["status"]["name"], issue["fields"]["priority"]["name"],issue["fields"]["project"]["name"],
        assignee_name, updated_at, Json(raw_data)
    ))
    task_id = cursor.fetchone()[0]
    add_task_version(task_id, issue, updated_at, raw_data, cursor)

    # Увеличиваем счетчик новых issues и обновляем статистику для проекта
    new_issues_count += 1
    if project_name not in projects_statistics:
        projects_statistics[project_name] = {'new': 1, 'updated': 0}
    else:
        projects_statistics[project_name]['new'] += 1

def update_task(task_id, issue, updated_at, raw_data, cursor):
    global updated_issues_count, projects_statistics
    assignee_name = issue["fields"].get("assignee", {}).get("displayName", "") if issue["fields"].get("assignee") else ""
    project_name = issue["fields"]["project"]["name"]
    cursor.execute("""
        UPDATE jira_tasks SET summary = %s, description = %s, status = %s, 
        priority = %s, project = %s, assignee = %s, updated_at = %s, raw_data = %s WHERE task_id = %s
    """, (
        issue["fields"]["summary"], issue["fields"].get("description"),
        issue["fields"]["status"]["name"], issue["fields"]["priority"]["name"], issue["fields"]["project"]["name"],
        assignee_name, updated_at, Json(raw_data), task_id
    ))
    add_task_version(task_id, issue, updated_at, raw_data, cursor)
    # Увеличиваем счетчик обновленных issues и обновляем статистику для проекта
    updated_issues_count += 1
    if project_name not in projects_statistics:
        projects_statistics[project_name] = {'new': 0, 'updated': 1}
    else:
        projects_statistics[project_name]['updated'] += 1

def add_task_version(task_id, issue, version_timestamp, raw_data, cursor):
    # Также применяем здесь проверку
    assignee_name = issue["fields"].get("assignee", {}).get("displayName", "") if issue["fields"].get("assignee") else ""
    cursor.execute("""
        INSERT INTO task_versions (task_id, jira_key, summary, description, status, priority, project, assignee, version_timestamp, raw_data)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    """, (
        task_id, issue["key"], issue["fields"]["summary"], issue["fields"].get("description"),
        issue["fields"]["status"]["name"], issue["fields"]["priority"]["name"], issue["fields"]["project"]["name"],
        assignee_name, version_timestamp, Json(raw_data)
    ))

def get_issues_count_by_jql(domain, email, api_token, jql_query):
    """
    Возвращает количество задач в Jira по заданному JQL запросу.
    
    Параметры:
    domain (str): Домен вашего инстанса Jira.
    email (str): Email пользователя для аутентификации.
    api_token (str): API токен пользователя.
    jql_query (str): JQL запрос для фильтрации задач.
    
    Возвращает:
    int: Общее количество задач, соответствующих JQL запросу.
    """
    
    api_url = f'{domain}/rest/api/3/search'
    auth = HTTPBasicAuth(email, api_token)
    headers = {"Accept": "application/json"}
    query = {'jql': jql_query, 'maxResults': 0}
    
    response = requests.get(api_url, headers=headers, params=query, auth=auth)
    
    if response.status_code == 200:
        data = response.json()
        return data.get('total', 0)
    else:
        print(f"Ошибка запроса: {response.status_code}")
        return 0

def fetch_and_update_custom_fields(JIRA_URL, JIRA_EMAIL, JIRA_API_TOKEN, DBNAME, DBUSER, DBPASSWORD, DBHOST):
    # Подключение к JIRA
    jira = JIRA(server=JIRA_URL, basic_auth=(JIRA_EMAIL, JIRA_API_TOKEN))
    
    # Получение метаданных всех полей из JIRA
    fields = jira.fields()
    
    # Фильтрация кастомных полей
    custom_fields = {field['id']: field['name'] for field in fields if field['custom']}
    
    # Подключение к базе данных
    conn = psycopg2.connect(dbname=DBNAME, user=DBUSER, password=DBPASSWORD, host=DBHOST)
    cursor = conn.cursor()
    
    # Добавление информации о кастомных полях в таблицу custom_fields в базе данных
    for field_id, field_name in custom_fields.items():
        # Проверка наличия поля в БД
        cursor.execute("SELECT field_id FROM custom_fields WHERE field_id = %s", (field_id,))
        if cursor.fetchone() is None:
            # Добавление нового кастомного поля
            cursor.execute("INSERT INTO custom_fields (field_id, field_name) VALUES (%s, %s)", (field_id, field_name))
        else:
            # Обновление имени поля, если оно уже есть в БД
            cursor.execute("UPDATE custom_fields SET field_name = %s WHERE field_id = %s", (field_name, field_id))

    # Фиксация изменений в БД и закрытие соединения
    conn.commit()
    cursor.close()
    conn.close()

# После завершения обработки всех issues, выводим итоговую статистику
def print_statistics():
    global new_issues_count, updated_issues_count, projects_statistics
    print(f"Total new issues added: {new_issues_count}")
    print(f"Total issues updated: {updated_issues_count}")
    for project, stats in projects_statistics.items():
        print(f"Project {project}: New - {stats['new']}, Updated - {stats['updated']}")

if __name__ == "__main__":
    try:
        start_time = datetime.now(timezone.utc)
        log_details = {
            'id_comp': socket.gethostname(),
            'start_time': start_time,
        }
        

        # Создание парсера
        parser = argparse.ArgumentParser(description="Cкрипт с аргументами командной строки.")
        # Добавление аргументов
        parser.add_argument('period_type', help="Задаём тип периода запроса.", nargs='?', default='w')
        # Разбор аргументов
        args = parser.parse_args()
        PERIOD_TYPE=args.period_type

        log_details.update({'script_params': str(args)})

        fetch_and_update_custom_fields(JIRA_URL, JIRA_EMAIL, JIRA_API_TOKEN, DBNAME, DBUSER, DBPASSWORD, DBHOST)
        
        # Получаем текущую дату и время
        current_date_time = datetime.now()
        # Форматируем дату в строку в формате YYYY-MM-DD
        END_DATE = current_date_time.strftime('%Y-%m-%d')
        START_DATE=get_relative_date(current_date_time, PERIOD_TYPE).strftime('%Y-%m-%d')

        # JQL query to fetch issues
        JQL_QUERY = build_jql_query(JIRA_PROJECTS,START_DATE,END_DATE)
        #JQL_QUERY = 'project in ("HTF", "TSM", "FTRM", "HD", "IYNO") and key=IYNO-5695'
        # Connect to JIRA
        jira = JIRA(server=JIRA_URL, basic_auth=(JIRA_EMAIL, JIRA_API_TOKEN))

        # Connect to PostgreSQL database
        conn = psycopg2.connect(dbname=DBNAME, user=DBUSER, password=DBPASSWORD, host=DBHOST)
        cursor = conn.cursor()

        # общее число issues
        issues_count = get_issues_count_by_jql(JIRA_URL, JIRA_EMAIL, JIRA_API_TOKEN, JQL_QUERY)

        # Fetching issues from JIRA
        start_at = 0
        max_results = 100
        # Initialize the progress bar
        with tqdm(total=issues_count, desc="Processing issues") as pbar:
            while True:
                issues = jira.search_issues(JQL_QUERY, startAt=start_at, maxResults=max_results, json_result=True)   
                if not issues['issues']:
                    break
                for issue in issues['issues']:
                    upsert_task(issue, cursor)
                start_at += max_results
                pbar.update(len(issues['issues']))

        print_statistics()

        end_time = datetime.now(timezone.utc)
        duration = end_time - start_time

        log_details.update({
            'end_time': end_time,
            'duration': duration,
            'log_message': "Script completed successfully",
            'jql': JQL_QUERY,
            'total_issues_fetched': issues_count,
            'total_issues_new': new_issues_count,
            'total_issues_updated': updated_issues_count
        })

        log_script_execution(conn, log_details)
    except Exception as e:
        end_time = datetime.now(timezone.utc)
        duration = end_time - start_time
        log_details.update({
            'end_time': end_time,
            'duration': duration,
            'log_level': 'ERROR',
            'log_message': str(e),
            'errors': str(e)
        })
        log_script_execution(conn, log_details)
        raise
    finally:
        # Commit changes and close connections
        conn.commit()

        export2parquet()

        cursor.close()
        conn.close()