Skip to content

Commit cf8caf0

Browse files
committed
feat: add dynamic insert query generation
1 parent f408fbf commit cf8caf0

File tree

5 files changed

+500
-152
lines changed

5 files changed

+500
-152
lines changed
Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
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+
operation = "GREATEST" if real_name == "_timestamp" else "COALESCE"
41+
42+
query_params_properties.append((real_name, operation))
43+
updateable_model_fields.append(field_name)
44+
updateable_db_fields.append(real_name)
45+
46+
updateable_db_fields_clauses = [
47+
f"""
48+
{updateable_field}"""
49+
for updateable_field in updateable_db_fields
50+
]
51+
52+
conflict_clauses = []
53+
for field, op in query_params_properties:
54+
conflict_clauses.append(
55+
f"""
56+
{field} = {op}({table_name}.{field}, EXCLUDED.{field})"""
57+
)
58+
59+
query = f"""
60+
INSERT INTO {table_name} ({','.join(updateable_db_fields_clauses)}
61+
)
62+
VALUES (
63+
{', '.join(['%s'] * len(updateable_db_fields))}
64+
)
65+
ON CONFLICT (id)
66+
DO UPDATE SET{','.join(conflict_clauses)};
67+
"""
68+
69+
var_insert_queries[table_name] = {}
70+
var_insert_queries[table_name][
71+
"updateable_model_fields"
72+
] = updateable_model_fields
73+
var_insert_queries[table_name]["query"] = query
74+
75+
# Read the template file
76+
template_path = os.path.join(
77+
os.path.dirname(__file__), "templates", "insert_queries.txt.j2"
78+
)
79+
with open(template_path, "r") as template_file:
80+
template_content = template_file.read()
81+
82+
# Render the template with the variables
83+
template = Template(template_content)
84+
rendered_content = template.render(
85+
timestamp=datetime.now(),
86+
checkouts=var_insert_queries["checkouts"],
87+
issues=var_insert_queries["issues"],
88+
builds=var_insert_queries["builds"],
89+
tests=var_insert_queries["tests"],
90+
incidents=var_insert_queries["incidents"],
91+
)
92+
93+
# Write the result to a Python file
94+
output_path = os.path.join(
95+
os.path.dirname(__file__), "generated", "insert_queries.py"
96+
)
97+
os.makedirs(os.path.dirname(output_path), exist_ok=True)
98+
with open(output_path, "w") as output_file:
99+
output_file.write(rendered_content)
100+
101+
self.stdout.write(
102+
self.style.SUCCESS(
103+
f"Successfully generated insert queries at {output_path}"
104+
)
105+
)

0 commit comments

Comments
 (0)