Source code for operators.linear_issues_to_s3
"""
For now, the scope of this operator is to read in the response from Linear API
and store the resulting data as parquet in S3.
Then CET can read the data from S3 and create raw table
(i.e., columns with subfields still in JSON format)
TODO: Include parameter to create transformed table with flattened columns
from the raw table, i.e., columns with subfields flattened out
"""
from typing import TYPE_CHECKING
from airflow.models import BaseOperator
from airflow.models.connection import Connection
from airflow.utils.context import Context
from airflow.utils.decorators import apply_defaults
from config import Config
from util.base import BaseUtil
if TYPE_CHECKING:
from typing import Any, Dict, List
# Constants
S3_ROOT = "s3://"
LAKE_BUCKET = Config.DATA_LAKE_BUCKET
LAKE_DB = Config.DATA_LAKE_DB
RAW = Config.RAW_PREFIX
PRODUCT_NAME = "linear"
SUBPRODUCT_NAME = "issues"
TABLE_NAME = "r_linear_issues"
[docs]
class LinearIssuesToS3Operator(BaseOperator):
"""
Custom operator to fetch issues from Linear API.
Store as parquet files in S3.
- fetches linear_api_key from Airflow connections
- s3_bucket_name is default
- s3_key is default -- partition data in S3 but when creating Athena table, use the latest data
(i.e., table in Athena will always point to the latest data in S3 -- overwritten)
"""
ui_color = "#FFFFFF"
@apply_defaults
def __init__(
self: object,
*args,
**kwargs,
):
super().__init__(
*args,
**kwargs,
)
[docs]
@staticmethod
def format_date(
date_str: str,
) -> str:
"""
Formate date string for readability if it exists
:param date_str: The date string to format
:type date_str: str
:return: Formatted date string or None if input is empty
:rtype: str or None
"""
from datetime import datetime
if not date_str:
return None
try:
dt = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
return dt.strftime("%Y-%m-%d %H:%M:%S")
except:
return date_str
[docs]
@staticmethod
def safe_get(
obj: object,
*keys: str,
default: object = None,
) -> object:
"""
Safely get a value from a nested dictionary
:param obj: The object to extract the value from
:rtype obj: object
:param keys: The keys to traverse the object
:type keys: str
:param default: The default value to return if the key is not found
:type default: object
:return: The value at the specified keys or the default value
:rtype: object
"""
if obj is None:
return default
current = obj
for key in keys:
try:
if isinstance(current, dict):
if key not in current:
return default
current = current[key]
else:
current = getattr(current, key)
if current is None:
return default
except (KeyError, AttributeError, TypeError):
return default
return current
[docs]
def format_issue(
issue: dict,
) -> dict:
"""
Format issue data for display with safe attribute access
:param issue: The issue data to format
:return: A dictionary with formatted issue data
:rtype: dict
"""
formatted_issue = {}
# Basic issue info
formatted_issue["ID"] = (
LinearIssuesToS3Operator.LinearIssuesToS3Operator.safe_get(issue, "id")
)
formatted_issue["Identifier"] = (
LinearIssuesToS3Operator.LinearIssuesToS3Operator.safe_get(
issue, "identifier"
)
)
formatted_issue["Number"] = (
LinearIssuesToS3Operator.LinearIssuesToS3Operator.safe_get(issue, "number")
)
formatted_issue["Title"] = LinearIssuesToS3Operator.safe_get(issue, "title")
formatted_issue["Description"] = LinearIssuesToS3Operator.safe_get(
issue, "description"
)
formatted_issue["Description State"] = LinearIssuesToS3Operator.safe_get(
issue, "descriptionState"
)
formatted_issue["URL"] = LinearIssuesToS3Operator.safe_get(issue, "url")
formatted_issue["Branch Name"] = LinearIssuesToS3Operator.safe_get(
issue, "branchName"
)
if formatted_issue["Branch Name"]:
formatted_issue["Branch Name"] = formatted_issue["Branch Name"].replace(
"susheel/", ""
)
# Metrics and organization
formatted_issue["Estimate"] = LinearIssuesToS3Operator.safe_get(
issue, "estimate"
)
formatted_issue["Priority"] = LinearIssuesToS3Operator.safe_get(
issue, "priority"
)
formatted_issue["Priority Label"] = LinearIssuesToS3Operator.safe_get(
issue, "priorityLabel"
)
formatted_issue["Sort Order"] = LinearIssuesToS3Operator.safe_get(
issue, "sortOrder"
)
formatted_issue["Board Order"] = LinearIssuesToS3Operator.safe_get(
issue, "boardOrder"
)
formatted_issue["Sub Issue Sort Order"] = LinearIssuesToS3Operator.safe_get(
issue, "subIssueSortOrder"
)
formatted_issue["Previous Identifiers"] = LinearIssuesToS3Operator.safe_get(
issue, "previousIdentifiers"
)
formatted_issue["Customer Ticket Count"] = LinearIssuesToS3Operator.safe_get(
issue, "customerTicketCount"
)
formatted_issue["Integration Source Type"] = LinearIssuesToS3Operator.safe_get(
issue, "integrationSourceType"
)
# Team info
team = LinearIssuesToS3Operator.safe_get(issue, "team")
if team:
formatted_issue["Team ID"] = LinearIssuesToS3Operator.safe_get(team, "id")
formatted_issue["Team"] = LinearIssuesToS3Operator.safe_get(team, "name")
formatted_issue["Team Key"] = LinearIssuesToS3Operator.safe_get(team, "key")
formatted_issue["Team Description"] = LinearIssuesToS3Operator.safe_get(
team, "description"
)
# State info
state = LinearIssuesToS3Operator.safe_get(issue, "state")
if state:
formatted_issue["Status ID"] = LinearIssuesToS3Operator.safe_get(
state, "id"
)
formatted_issue["Status"] = LinearIssuesToS3Operator.safe_get(state, "name")
formatted_issue["Status Type"] = LinearIssuesToS3Operator.safe_get(
state, "type"
)
formatted_issue["Status Color"] = LinearIssuesToS3Operator.safe_get(
state, "color"
)
formatted_issue["Status Description"] = LinearIssuesToS3Operator.safe_get(
state, "description"
)
# Project info
project = LinearIssuesToS3Operator.safe_get(issue, "project")
if project:
formatted_issue["Project ID"] = LinearIssuesToS3Operator.safe_get(
project, "id"
)
formatted_issue["Project"] = LinearIssuesToS3Operator.safe_get(
project, "name"
)
formatted_issue["Project Description"] = LinearIssuesToS3Operator.safe_get(
project, "description"
)
formatted_issue["Project Icon"] = LinearIssuesToS3Operator.safe_get(
project, "icon"
)
formatted_issue["Project Color"] = LinearIssuesToS3Operator.safe_get(
project, "color"
)
formatted_issue["Project State"] = LinearIssuesToS3Operator.safe_get(
project, "state"
)
formatted_issue["Project Slug ID"] = LinearIssuesToS3Operator.safe_get(
project, "slugId"
)
formatted_issue["Project Start Date"] = (
LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(project, "startDate")
)
)
formatted_issue["Project Target Date"] = (
LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(project, "targetDate")
)
)
# Creator info
creator = LinearIssuesToS3Operator.safe_get(issue, "creator")
if creator:
formatted_issue["Creator ID"] = LinearIssuesToS3Operator.safe_get(
creator, "id"
)
formatted_issue["Creator"] = LinearIssuesToS3Operator.safe_get(
creator, "name"
)
formatted_issue["Creator Email"] = LinearIssuesToS3Operator.safe_get(
creator, "email"
)
formatted_issue["Creator Display Name"] = LinearIssuesToS3Operator.safe_get(
creator, "displayName"
)
# Assignee info
assignee = LinearIssuesToS3Operator.safe_get(issue, "assignee")
if assignee:
formatted_issue["Assignee ID"] = LinearIssuesToS3Operator.safe_get(
assignee, "id"
)
formatted_issue["Assignee"] = LinearIssuesToS3Operator.safe_get(
assignee, "name"
)
formatted_issue["Assignee Email"] = LinearIssuesToS3Operator.safe_get(
assignee, "email"
)
formatted_issue["Assignee Display Name"] = (
LinearIssuesToS3Operator.safe_get(assignee, "displayName")
)
# Subscribers
subscribers_nodes = LinearIssuesToS3Operator.safe_get(
issue, "subscribers", "nodes", default=[]
)
subscriber_ids = []
subscriber_names = []
for subscriber in subscribers_nodes:
if subscriber:
subscriber_id = LinearIssuesToS3Operator.safe_get(subscriber, "id")
subscriber_name = LinearIssuesToS3Operator.safe_get(subscriber, "name")
if subscriber_id:
subscriber_ids.append(subscriber_id)
if subscriber_name:
subscriber_names.append(subscriber_name)
formatted_issue["Subscriber IDs"] = subscriber_ids if subscriber_ids else None
formatted_issue["Subscribers"] = subscriber_names if subscriber_names else None
# Labels
labels_nodes = LinearIssuesToS3Operator.safe_get(
issue, "labels", "nodes", default=[]
)
label_ids = []
label_names = []
for label in labels_nodes:
if label:
label_id = LinearIssuesToS3Operator.safe_get(label, "id")
label_name = LinearIssuesToS3Operator.safe_get(label, "name")
if label_id:
label_ids.append(label_id)
if label_name:
label_names.append(label_name)
formatted_issue["Label IDs"] = label_ids if label_ids else None
formatted_issue["Labels"] = label_names if label_names else None
# Comments - Store as JSON
comments_nodes = LinearIssuesToS3Operator.safe_get(
issue, "comments", "nodes", default=[]
)
comments = [
{
"id": LinearIssuesToS3Operator.safe_get(comment, "id"),
"body": LinearIssuesToS3Operator.safe_get(comment, "body"),
"user": LinearIssuesToS3Operator.safe_get(comment, "user", "name"),
"date": LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(comment, "createdAt")
),
}
for comment in comments_nodes
if comment
]
formatted_issue["Comments"] = comments
# Children - Store as JSON
children_nodes = LinearIssuesToS3Operator.safe_get(
issue, "children", "nodes", default=[]
)
children = [
{
"id": LinearIssuesToS3Operator.safe_get(child, "id"),
"identifier": LinearIssuesToS3Operator.safe_get(child, "identifier"),
"title": LinearIssuesToS3Operator.safe_get(child, "title"),
"state": LinearIssuesToS3Operator.safe_get(child, "state", "name"),
}
for child in children_nodes
if child
]
formatted_issue["Children"] = children
# History - Store as JSON
history_nodes = LinearIssuesToS3Operator.safe_get(
issue, "history", "nodes", default=[]
)
history_entries = [
{
"id": LinearIssuesToS3Operator.safe_get(history, "id"),
"date": LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(history, "createdAt")
),
"actor": LinearIssuesToS3Operator.safe_get(history, "actor", "name"),
"change": (
f"{LinearIssuesToS3Operator.safe_get(history, 'fromState', 'name')} → {LinearIssuesToS3Operator.safe_get(history, 'toState', 'name')}"
if LinearIssuesToS3Operator.safe_get(history, "fromState", "name")
and LinearIssuesToS3Operator.safe_get(history, "toState", "name")
else ""
),
}
for history in history_nodes
if history
]
formatted_issue["History"] = history_entries
# Attachments - Store as JSON
attachments_nodes = LinearIssuesToS3Operator.safe_get(
issue, "attachments", "nodes", default=[]
)
attachments = [
{
"id": LinearIssuesToS3Operator.safe_get(attachment, "id"),
"title": LinearIssuesToS3Operator.safe_get(attachment, "title"),
"url": LinearIssuesToS3Operator.safe_get(attachment, "url"),
}
for attachment in attachments_nodes
if attachment
]
formatted_issue["Attachments"] = attachments
# Relations - Store as JSON
relations_nodes = LinearIssuesToS3Operator.safe_get(
issue, "relations", "nodes", default=[]
)
relations = [
{
"id": LinearIssuesToS3Operator.safe_get(relation, "id"),
"type": LinearIssuesToS3Operator.safe_get(relation, "type"),
"related_issue": {
"id": LinearIssuesToS3Operator.safe_get(
LinearIssuesToS3Operator.safe_get(relation, "relatedIssue"),
"id",
),
"identifier": LinearIssuesToS3Operator.safe_get(
LinearIssuesToS3Operator.safe_get(relation, "relatedIssue"),
"identifier",
),
},
}
for relation in relations_nodes
if relation
]
formatted_issue["Relations"] = relations
# Cycle info
cycle = LinearIssuesToS3Operator.safe_get(issue, "cycle")
if cycle:
formatted_issue["Cycle ID"] = LinearIssuesToS3Operator.safe_get(cycle, "id")
formatted_issue["Cycle Number"] = LinearIssuesToS3Operator.safe_get(
cycle, "number"
)
formatted_issue["Cycle Name"] = LinearIssuesToS3Operator.safe_get(
cycle, "name"
)
formatted_issue["Cycle Description"] = LinearIssuesToS3Operator.safe_get(
cycle, "description"
)
formatted_issue["Cycle Start"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(cycle, "startsAt")
)
formatted_issue["Cycle End"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(cycle, "endsAt")
)
formatted_issue["Cycle Completed"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(cycle, "completedAt")
)
formatted_issue["Cycle Progress"] = LinearIssuesToS3Operator.safe_get(
cycle, "progress"
)
# Parent issue
parent = LinearIssuesToS3Operator.safe_get(issue, "parent")
if parent:
formatted_issue["Parent ID"] = LinearIssuesToS3Operator.safe_get(
parent, "id"
)
formatted_issue["Parent Identifier"] = LinearIssuesToS3Operator.safe_get(
parent, "identifier"
)
formatted_issue["Parent Title"] = LinearIssuesToS3Operator.safe_get(
parent, "title"
)
formatted_issue["Parent Issue"] = (
f"{LinearIssuesToS3Operator.safe_get(parent, 'identifier')} - {LinearIssuesToS3Operator.safe_get(parent, 'title')}"
)
# Project milestone
milestone = LinearIssuesToS3Operator.safe_get(issue, "projectMilestone")
if milestone:
formatted_issue["Project Milestone ID"] = LinearIssuesToS3Operator.safe_get(
milestone, "id"
)
formatted_issue["Project Milestone"] = LinearIssuesToS3Operator.safe_get(
milestone, "name"
)
formatted_issue["Project Milestone Description"] = (
LinearIssuesToS3Operator.safe_get(milestone, "description")
)
formatted_issue["Project Milestone Target Date"] = (
LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(milestone, "targetDate")
)
)
# All timestamps
formatted_issue["Created"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "createdAt")
)
formatted_issue["Updated"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "updatedAt")
)
formatted_issue["Started"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "startedAt")
)
formatted_issue["Triaged"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "triagedAt")
)
formatted_issue["Completed"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "completedAt")
)
formatted_issue["Canceled"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "canceledAt")
)
formatted_issue["Archived"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "archivedAt")
)
formatted_issue["Auto Closed"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "autoClosedAt")
)
formatted_issue["Auto Archived"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "autoArchivedAt")
)
formatted_issue["Due Date"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "dueDate")
)
formatted_issue["Snoozed Until"] = LinearIssuesToS3Operator.format_date(
LinearIssuesToS3Operator.safe_get(issue, "snoozedUntilAt")
)
return formatted_issue
[docs]
@staticmethod
def get_all_linear_issues(
api_key: str,
) -> "List[Dict[str, Any]]":
"""
Fetch ALL Linear issues with pagination and retry logic.
:param api_key: Linear API key
:type api_key: str
:return: List of dictionaries containing all issues
:rtype: List[Dict[str, Any]]
"""
import time
# GraphQL query with pagination parameters
query = """
query GetDetailedIssues($first: Int, $after: String) {
issues(first: $first, after: $after) {
nodes {
id
identifier
title
description
descriptionState
branchName
estimate
priority
priorityLabel
sortOrder
autoClosedAt
autoArchivedAt
boardOrder
subIssueSortOrder
previousIdentifiers
customerTicketCount
number
createdAt
updatedAt
archivedAt
startedAt
completedAt
canceledAt
dueDate
triagedAt
snoozedUntilAt
integrationSourceType
url
team {
id
name
key
description
}
state {
id
name
color
description
type
}
project {
id
name
description
icon
color
state
slugId
startDate
targetDate
}
creator {
id
name
email
displayName
avatarUrl
}
assignee {
id
name
email
displayName
avatarUrl
}
subscribers {
nodes {
id
name
displayName
}
}
labels {
nodes {
id
name
color
description
}
}
comments {
nodes {
id
body
createdAt
updatedAt
user {
id
name
email
}
}
}
cycle {
id
number
name
description
startsAt
endsAt
completedAt
progress
}
parent {
id
identifier
title
}
children {
nodes {
id
identifier
title
state {
name
type
}
}
}
projectMilestone {
id
name
description
targetDate
}
history {
nodes {
id
createdAt
actor {
id
name
}
fromState {
id
name
}
toState {
id
name
}
}
}
attachments {
nodes {
id
title
url
createdAt
}
}
relations {
nodes {
id
type
relatedIssue {
id
identifier
title
}
}
}
}
pageInfo {
hasNextPage
endCursor
}
}
}
"""
# Linear GraphQL API endpoint
url = "https://api.linear.app/graphql"
# Headers with authorization
headers = {"Authorization": api_key, "Content-Type": "application/json"}
all_issues = []
has_next_page = True
after_cursor = None
page_count = 0
page_size = 50 # Reduced from 100 to lower the risk of timeouts
# Loop through all pages of results
while has_next_page:
page_count += 1
variables = {
"first": page_size,
"after": after_cursor,
}
try:
# Use BaseUtil.post_request for the API call
response = BaseUtil.post_request(
url=url,
headers=headers,
data={
"query": query,
"variables": variables,
},
timeout=30,
retry_count=5,
retry_sleep=3,
)
# Parse the response
data = response.json()
# Check for GraphQL errors
if "errors" in data:
raise Exception(f"GraphQL errors: {data['errors']}")
# Extract the issues and page info
issues_data = data.get("data", {}).get("issues", {})
issues = issues_data.get("nodes", [])
page_info = issues_data.get("pageInfo", {})
# Add issues from this page to our collection
all_issues.extend(issues)
# Check if there are more pages
has_next_page = page_info.get("hasNextPage", False)
after_cursor = page_info.get("endCursor", None)
# Print progress
print(
f"Fetched page {page_count} with {len(issues)} issues. Total so far: {len(all_issues)}"
)
except Exception as e:
print(f"Error fetching page {page_count}: {e}")
raise
# Add a delay between successful requests to avoid rate limits
if has_next_page:
time.sleep(1.5) # Increased from 0.5 to be more conservative
print(f"Completed! Total issues fetched: {len(all_issues)}")
return all_issues
[docs]
@staticmethod
def write_to_s3(
issues: "List[Dict[str, Any]]",
) -> None:
"""
Write the list of issues to S3 as a parquet file.
Create table in Athena with the latest data.
:param issues: List of issue dictionaries to write to S3
:type issues: List[Dict[str, Any]]
:return: None
:rtype: None
"""
import awswrangler as wr
import pandas as pd
if not issues:
print("No issues to write to S3.")
return
wr.s3.to_parquet(
df=pd.DataFrame(issues),
path=f"{S3_ROOT}{LAKE_BUCKET}/{RAW}/{PRODUCT_NAME}/{SUBPRODUCT_NAME}",
dataset=True,
mode="overwrite",
database=LAKE_DB,
table=TABLE_NAME,
dtype={key: "string" for key in issues[0].keys()},
)
[docs]
def execute(
self,
context: Context,
) -> None:
"""
Get the Linear API key from Airflow connections
"""
import json
LINEAR_API_KEY = json.loads(BaseUtil.get_aws_secret(Config.LINEAR_API_KEY))["api_key"]
try:
# Get API key from environment variable for security
if not LINEAR_API_KEY:
raise ValueError("LINEAR_API_KEY environment variable not set")
print("Fetching ALL Linear issues with pagination...")
# Get all issues using the pagination function
issues = []
batch_size = 500 # Process issues in batches
try:
issues = LinearIssuesToS3Operator.get_all_linear_issues(LINEAR_API_KEY)
except Exception as e:
print(f"Error fetching all issues: {e}")
print("Attempting to save any issues fetched before the error...")
print(f"\nFound {len(issues)} total issues")
if not issues:
print(
"No issues were fetched. Please check your API key and connection."
)
return
# Write issues to S3
print("Writing issues to S3...")
LinearIssuesToS3Operator.write_to_s3(issues)
print("Issues written to S3 successfully.")
except Exception as e:
print(f"Error: {e}")
import traceback
traceback.print_exc()