85 lines
3.7 KiB
Python
85 lines
3.7 KiB
Python
from django.core.management.base import BaseCommand
|
|
from django.apps import apps
|
|
from django.db import connection, transaction
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Synchronizes the geojson field (and lon/lat if missing) from geom for all assets/models using PostGIS SQL."
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument(
|
|
'--dry-run',
|
|
action='store_true',
|
|
help='Perform a dry run without modifying the database.',
|
|
)
|
|
parser.add_argument(
|
|
'--only-missing',
|
|
action='store_true',
|
|
help='Only update records where geojson is null or empty.',
|
|
)
|
|
|
|
def handle(self, *args, **options):
|
|
dry_run = options['dry_run']
|
|
only_missing = options['only_missing']
|
|
|
|
total_scanned = 0
|
|
total_updated = 0
|
|
|
|
with connection.cursor() as cursor:
|
|
for model in apps.get_models():
|
|
fields = [f.name for f in model._meta.get_fields()]
|
|
if 'geom' not in fields or 'geojson' not in fields:
|
|
continue
|
|
|
|
table = model._meta.db_table
|
|
is_json_field = model._meta.get_field('geojson').get_internal_type() == 'JSONField'
|
|
has_lon = 'lon' in fields
|
|
has_lat = 'lat' in fields
|
|
|
|
# Count total candidates with non-null geom
|
|
cursor.execute(f'SELECT COUNT(*) FROM "{table}" WHERE geom IS NOT NULL;')
|
|
count = cursor.fetchone()[0]
|
|
if count == 0:
|
|
continue
|
|
|
|
self.stdout.write(f"Processing {model._meta.app_label}.{model.__name__} ({table}, {count} candidate records)...")
|
|
|
|
if only_missing:
|
|
where_clause = "geom IS NOT NULL AND (geojson IS NULL OR geojson::text = '' OR geojson::text = 'null')"
|
|
else:
|
|
where_clause = "geom IS NOT NULL"
|
|
|
|
cursor.execute(f'SELECT COUNT(*) FROM "{table}" WHERE {where_clause};')
|
|
records_to_update = cursor.fetchone()[0]
|
|
|
|
if records_to_update == 0:
|
|
self.stdout.write(self.style.SUCCESS(f" {model.__name__}: 0 records need update."))
|
|
continue
|
|
|
|
if dry_run:
|
|
self.stdout.write(self.style.SUCCESS(f" [DRY-RUN] {model.__name__}: Would update {records_to_update}/{count} records."))
|
|
total_scanned += count
|
|
total_updated += records_to_update
|
|
continue
|
|
|
|
cast_geojson = "ST_AsGeoJSON(ST_Transform(geom, 4326))::jsonb" if is_json_field else "ST_AsGeoJSON(ST_Transform(geom, 4326))"
|
|
|
|
set_clauses = [f"geojson = {cast_geojson}"]
|
|
if has_lon:
|
|
set_clauses.append("lon = COALESCE(lon, ST_X(ST_PointOnSurface(ST_Transform(geom, 4326))))")
|
|
if has_lat:
|
|
set_clauses.append("lat = COALESCE(lat, ST_Y(ST_PointOnSurface(ST_Transform(geom, 4326))))")
|
|
|
|
set_str = ", ".join(set_clauses)
|
|
update_sql = f'UPDATE "{table}" SET {set_str} WHERE {where_clause};'
|
|
|
|
with transaction.atomic():
|
|
cursor.execute(update_sql)
|
|
updated_count = cursor.rowcount
|
|
|
|
self.stdout.write(self.style.SUCCESS(f" {model.__name__}: {updated_count}/{count} records updated in DB."))
|
|
total_scanned += count
|
|
total_updated += updated_count
|
|
|
|
prefix = "[DRY-RUN] Would update" if dry_run else "Successfully updated"
|
|
self.stdout.write(self.style.SUCCESS(f"\n{prefix} {total_updated} / {total_scanned} records across all models."))
|