Skip to content

Commit 49cec20

Browse files
committed
feat: change ingester data insertion policy
Now allows for overwrites of null data Also adds a command to generate the insertion queries dynamically for any model Closes #1552
1 parent 1b6daa3 commit 49cec20

File tree

6 files changed

+518
-29
lines changed

6 files changed

+518
-29
lines changed
Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
from django.core.management.base import BaseCommand
2+
from kernelCI_app.typeModels.modelTypes import MODEL_MAP
3+
import os
4+
from datetime import datetime
5+
from jinja2 import Template
6+
7+
8+
class Command(BaseCommand):
9+
help = """
10+
Dynamically generates the insert queries for all models, storing them in a specific file.
11+
Gives priority to the existing data in the database.
12+
13+
This command should not executed in runtime.
14+
"""
15+
16+
def add_arguments(self, parser):
17+
# Add custom command arguments here
18+
pass
19+
20+
def handle(self, *args, **options):
21+
var_insert_queries = {}
22+
23+
for table_name in MODEL_MAP.keys():
24+
model = MODEL_MAP[table_name]
25+
26+
updateable_model_fields: list[str] = []
27+
updateable_db_fields: list[str] = []
28+
query_params_properties: list[tuple[str, str]] = []
29+
30+
for field in model._meta.fields:
31+
if field.generated:
32+
continue
33+
34+
field_name = (
35+
field.name + "_id"
36+
if field.get_internal_type() == "ForeignKey"
37+
else field.name
38+
)
39+
real_name = field.db_column or field_name
40+
updateable_model_fields.append(field_name)
41+
updateable_db_fields.append(real_name)
42+
43+
operation = "GREATEST" if real_name == "_timestamp" else "COALESCE"
44+
# Fields that are never null don't need to have conflict clauses (except _timestamp)
45+
if field.null or real_name == "_timestamp":
46+
query_params_properties.append((real_name, operation))
47+
48+
updateable_db_fields_clauses = [
49+
f"""
50+
{updateable_field}"""
51+
for updateable_field in updateable_db_fields
52+
]
53+
54+
conflict_clauses = []
55+
for field, op in query_params_properties:
56+
conflict_clauses.append(
57+
f"""
58+
{field} = {op}({table_name}.{field}, EXCLUDED.{field})"""
59+
)
60+
61+
query = f"""
62+
INSERT INTO {table_name} ({','.join(updateable_db_fields_clauses)}
63+
)
64+
VALUES (
65+
{', '.join(['%s'] * len(updateable_db_fields))}
66+
)
67+
ON CONFLICT (id)
68+
DO UPDATE SET{','.join(conflict_clauses)};
69+
"""
70+
71+
var_insert_queries[table_name] = {}
72+
var_insert_queries[table_name][
73+
"updateable_model_fields"
74+
] = updateable_model_fields
75+
var_insert_queries[table_name]["query"] = query
76+
77+
# Read the template file
78+
template_path = os.path.join(
79+
os.path.dirname(__file__), "templates", "insert_queries.txt.j2"
80+
)
81+
with open(template_path, "r") as template_file:
82+
template_content = template_file.read()
83+
84+
# Render the template with the variables
85+
template = Template(template_content)
86+
rendered_content = template.render(
87+
timestamp=datetime.now(),
88+
checkouts=var_insert_queries["checkouts"],
89+
issues=var_insert_queries["issues"],
90+
builds=var_insert_queries["builds"],
91+
tests=var_insert_queries["tests"],
92+
incidents=var_insert_queries["incidents"],
93+
)
94+
95+
# Write the result to a Python file
96+
output_path = os.path.join(
97+
os.path.dirname(__file__), "generated", "insert_queries.py"
98+
)
99+
os.makedirs(os.path.dirname(output_path), exist_ok=True)
100+
with open(output_path, "w") as output_file:
101+
output_file.write(rendered_content)
102+
103+
self.stdout.write(
104+
self.style.SUCCESS(
105+
f"Successfully generated insert queries at {output_path}"
106+
)
107+
)

0 commit comments

Comments
 (0)