Skip to content

Commit

Permalink
psycopg2 cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@axelstudios
axelstudios committed Jan 17, 2025
1 parent 31d08fd commit c179d2d
Show file tree
Hide file tree
Showing 6 changed files with 77 additions and 78 deletions.
4 changes: 2 additions & 2 deletions config/settings/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@
SEED Platform (TM), Copyright (c) Alliance for Sustainable Energy, LLC, and other contributors.
See also https://github.com/SEED-platform/seed/blob/main/LICENSE.md
"""
from __future__ import annotations

import os
from typing import Union

from django.utils.translation import gettext_lazy as _
from kombu.serialization import register
Expand Down Expand Up @@ -307,7 +307,7 @@
}


def yn(s: bool | str) -> bool:
def yn(s: Union[bool, str]) -> bool:
if isinstance(s, bool):
return s
if isinstance(s, str):
Expand Down
2 changes: 1 addition & 1 deletion config/settings/test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,4 +84,4 @@
SF_SECURITY_TOKEN = os.environ.get("SF_SECURITY_TOKEN", "")

# load small EEEJ dataset for testing
EEEJ_LOAD_SMALL_TEST_DATASET = yn(os.environ.get("EEEJ_LOAD_SMALL_TEST_DATASET", "True"))
EEEJ_LOAD_SMALL_TEST_DATASET = yn(os.environ.get("EEEJ_LOAD_SMALL_TEST_DATASET", "True")) # noqa: F405
2 changes: 1 addition & 1 deletion requirements/base.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
django==4.2.18

# Persistence stores
psycopg[binary]==3.2.4
psycopg2-binary==2.9.10

# Background process management
celery==5.4.0
Expand Down
111 changes: 55 additions & 56 deletions seed/migrations/0242_rename_indices.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,100 +4,99 @@


class Migration(migrations.Migration):

dependencies = [
('seed', '0241_alter_meter_type'),
("seed", "0241_alter_meter_type"),
]

operations = [
migrations.RenameIndex(
model_name='note',
new_name='seed_note_organiz_909e1a_idx',
old_fields=('organization', 'note_type'),
model_name="note",
new_name="seed_note_organiz_909e1a_idx",
old_fields=("organization", "note_type"),
),
migrations.RenameIndex(
model_name='propertyauditlog',
new_name='seed_proper_parent__8a8a20_idx',
old_fields=('parent_state1', 'parent_state2'),
model_name="propertyauditlog",
new_name="seed_proper_parent__8a8a20_idx",
old_fields=("parent_state1", "parent_state2"),
),
migrations.RenameIndex(
model_name='propertyauditlog',
new_name='seed_proper_state_i_9c4415_idx',
old_fields=('state', 'name'),
model_name="propertyauditlog",
new_name="seed_proper_state_i_9c4415_idx",
old_fields=("state", "name"),
),
migrations.RenameIndex(
model_name='propertymeasure',
new_name='seed_proper_propert_3e8b15_idx',
old_fields=('property_measure_name', 'property_state'),
model_name="propertymeasure",
new_name="seed_proper_propert_3e8b15_idx",
old_fields=("property_measure_name", "property_state"),
),
migrations.RenameIndex(
model_name='propertystate',
new_name='seed_proper_import__302589_idx',
old_fields=('import_file', 'data_state'),
model_name="propertystate",
new_name="seed_proper_import__302589_idx",
old_fields=("import_file", "data_state"),
),
migrations.RenameIndex(
model_name='propertystate',
new_name='seed_proper_import__a13e5f_idx',
old_fields=('import_file', 'data_state', 'merge_state'),
model_name="propertystate",
new_name="seed_proper_import__a13e5f_idx",
old_fields=("import_file", "data_state", "merge_state"),
),
migrations.RenameIndex(
model_name='propertystate',
new_name='seed_proper_import__ddad37_idx',
old_fields=('import_file', 'data_state', 'source_type'),
model_name="propertystate",
new_name="seed_proper_import__ddad37_idx",
old_fields=("import_file", "data_state", "source_type"),
),
migrations.RenameIndex(
model_name='propertystate',
new_name='seed_proper_hash_ob_6497e4_idx',
old_fields=('hash_object',),
model_name="propertystate",
new_name="seed_proper_hash_ob_6497e4_idx",
old_fields=("hash_object",),
),
migrations.RenameIndex(
model_name='propertyview',
new_name='seed_proper_state_i_dae98d_idx',
old_fields=('state', 'cycle'),
model_name="propertyview",
new_name="seed_proper_state_i_dae98d_idx",
old_fields=("state", "cycle"),
),
migrations.RenameIndex(
model_name='taxlotauditlog',
new_name='seed_taxlot_parent__a3e767_idx',
old_fields=('parent_state1', 'parent_state2'),
model_name="taxlotauditlog",
new_name="seed_taxlot_parent__a3e767_idx",
old_fields=("parent_state1", "parent_state2"),
),
migrations.RenameIndex(
model_name='taxlotauditlog',
new_name='seed_taxlot_state_i_093659_idx',
old_fields=('state', 'name'),
model_name="taxlotauditlog",
new_name="seed_taxlot_state_i_093659_idx",
old_fields=("state", "name"),
),
migrations.RenameIndex(
model_name='taxlotproperty',
new_name='seed_taxlot_propert_be033b_idx',
old_fields=('property_view', 'taxlot_view'),
model_name="taxlotproperty",
new_name="seed_taxlot_propert_be033b_idx",
old_fields=("property_view", "taxlot_view"),
),
migrations.RenameIndex(
model_name='taxlotproperty',
new_name='seed_taxlot_cycle_i_b3c514_idx',
old_fields=('cycle', 'taxlot_view'),
model_name="taxlotproperty",
new_name="seed_taxlot_cycle_i_b3c514_idx",
old_fields=("cycle", "taxlot_view"),
),
migrations.RenameIndex(
model_name='taxlotproperty',
new_name='seed_taxlot_cycle_i_577466_idx',
old_fields=('cycle', 'property_view'),
model_name="taxlotproperty",
new_name="seed_taxlot_cycle_i_577466_idx",
old_fields=("cycle", "property_view"),
),
migrations.RenameIndex(
model_name='taxlotstate',
new_name='seed_taxlot_import__8f600b_idx',
old_fields=('import_file', 'data_state'),
model_name="taxlotstate",
new_name="seed_taxlot_import__8f600b_idx",
old_fields=("import_file", "data_state"),
),
migrations.RenameIndex(
model_name='taxlotstate',
new_name='seed_taxlot_hash_ob_6c9805_idx',
old_fields=('hash_object',),
model_name="taxlotstate",
new_name="seed_taxlot_hash_ob_6c9805_idx",
old_fields=("hash_object",),
),
migrations.RenameIndex(
model_name='taxlotstate',
new_name='seed_taxlot_import__bc9aae_idx',
old_fields=('import_file', 'data_state', 'merge_state'),
model_name="taxlotstate",
new_name="seed_taxlot_import__bc9aae_idx",
old_fields=("import_file", "data_state", "merge_state"),
),
migrations.RenameIndex(
model_name='taxlotview',
new_name='seed_taxlot_state_i_aaa042_idx',
old_fields=('state', 'cycle'),
model_name="taxlotview",
new_name="seed_taxlot_state_i_aaa042_idx",
old_fields=("state", "cycle"),
),
]
24 changes: 15 additions & 9 deletions seed/models/meters.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
from django.db import IntegrityError, connection, models
from django.db.models.signals import pre_save
from django.dispatch import receiver
from psycopg2.extras import execute_values

from seed.models import Property, Scenario
from seed.models.inventory_groups import InventoryGroupMapping
Expand Down Expand Up @@ -177,19 +176,26 @@ def copy_readings(self, source_meter, overlaps_possible=True):
if overlaps_possible:
sql = (
"INSERT INTO seed_meterreading(meter_id, start_time, end_time, reading, source_unit, conversion_factor) "
"VALUES %s "
"VALUES (%s, %s, %s, %s, %s, %s) "
"ON CONFLICT (meter_id, start_time, end_time) "
"DO UPDATE SET reading=excluded.reading, source_unit=excluded.source_unit, conversion_factor=excluded.conversion_factor "
"DO UPDATE SET reading=excluded.reading, source_unit = excluded.source_unit, conversion_factor = excluded.conversion_factor "
"RETURNING reading"
)

with connection.cursor() as cursor:
execute_values(
cursor,
sql,
source_meter.meter_readings.values(),
template=f"({self.id}, %(start_time)s, %(end_time)s, %(reading)s, %(source_unit)s, %(conversion_factor)s)",
data = [
(
self.id,
reading["start_time"],
reading["end_time"],
reading["reading"],
reading["source_unit"],
reading["conversion_factor"],
)
for reading in source_meter.meter_readings.values()
]

with connection.cursor() as cursor:
cursor.executemany(sql, data)
else:
readings = {
MeterReading(
Expand Down
12 changes: 3 additions & 9 deletions seed/serializers/meter_readings.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
from django.core.exceptions import ValidationError
from django.db import connection
from django.utils.timezone import make_aware
from psycopg2.extras import execute_values
from pytz import timezone
from rest_framework import serializers

Expand All @@ -32,20 +31,15 @@ def to_internal_value(self, data):
def create(self, validated_data) -> list[MeterReading]:
upsert_sql = (
f"INSERT INTO seed_meterreading({', '.join(meter_fields)}) " # noqa: S608
"VALUES %s "
"VALUES (%(meter_id)s, %(start_time)s, %(end_time)s, %(reading)s, %(source_unit)s, %(conversion_factor)s) "
"ON CONFLICT (meter_id, start_time, end_time) "
"DO UPDATE SET reading=excluded.reading, source_unit=excluded.source_unit, conversion_factor=excluded.conversion_factor "
f"RETURNING {', '.join(meter_fields)}"
)

with connection.cursor() as cursor:
results: list[tuple] = execute_values(
cursor,
upsert_sql,
validated_data,
template="(%(meter_id)s, %(start_time)s, %(end_time)s, %(reading)s, %(source_unit)s, %(conversion_factor)s)",
fetch=True,
)
cursor.executemany(upsert_sql, validated_data)
results = cursor.fetchall()

# Convert list of tuples to list of MeterReadings for response
updated_readings = [MeterReading(**{field: result[i] for i, field in enumerate(meter_fields)}) for result in results]
Expand Down

0 comments on commit c179d2d

Please sign in to comment.