feat: speed up data insertion

dev.Dockerfile

@@ -50,5 +50,9 @@ RUN /app/setup-db.sh
 # ===============================
 FROM postgres:16.4
 
-# Copy preloaded, vacuumed database
-COPY --from=builder /var/lib/postgresql/data /var/lib/postgresql/data
+USER root
+RUN mkdir -p /var/lib/postgresql/data && chown -R postgres:postgres /var/lib/postgresql
+
+COPY --from=builder --chown=postgres:postgres /var/lib/postgresql/data /var/lib/postgresql/data
+
+USER postgres

scripts/export-and-load-autocomplete.sh

@@ -0,0 +1,78 @@
+#!/bin/bash
+
+set -euo pipefail
+
+# shellcheck disable=SC1091
+source ./scripts/util.sh
+
+# Function to export unique values to a CSV or SQL file
+export_unique_values() {
+    local query="$1"
+    local output_file="$2"
+    local format="${3:-csv}" # Default format is CSV
+    local base_name
+    base_name=$(basename "$output_file")
+
+    mkdir -p "$(dirname "$output_file")"
+
+    if [ "$format" == "csv" ]; then
+        log_info "Exporting unique values for the $base_name table in CSV format..."
+
+        psql -d postgres -c "\copy ($query) TO '$output_file' CSV HEADER;"
+        log_success "Exported CSV file saved to $output_file"
+
+    elif [ "$format" == "sql" ]; then
+        log_info "Exporting unique values for the $base_name table in SQL format..."
+
+        local temp_table="temp_export"
+        local output_sql="$output_file.sql"
+
+        psql -d postgres -c "
+            DROP TABLE IF EXISTS $temp_table;
+            CREATE TABLE $temp_table AS
+            SELECT row_number() OVER () AS id, * FROM ($query) AS subquery;
+        "
+
+        pg_dump -U postgres --data-only --table="$temp_table" postgres >"$output_sql"
+        psql -d postgres -c "DROP TABLE IF EXISTS $temp_table;"
+
+        log_success "Exported SQL file saved to $output_sql"
+    else
+        log_error "Invalid format specified. Use 'csv' or 'sql'."
+        exit 1
+    fi
+}
+
+# Transfers data from a CSV file to the database.
+transfer_csv_to_database() {
+    local table_name="$1"
+    local csv_file_path="$2"
+    local columns="$3"
+    local delimiter="$4"
+
+    if [ ! -f "$csv_file_path" ]; then
+        log_error "The CSV file '$csv_file_path' doesn't exist."
+        exit 1
+    fi
+
+    log_info "Copying data from '$csv_file_path' to the database table '$table_name'..."
+
+    psql -d postgres -c "\copy $table_name($columns) FROM '$csv_file_path' DELIMITER '$delimiter' CSV HEADER;"
+
+    log_success "Transfer of '$csv_file_path' to the database table '$table_name' completed successfully."
+}
+
+# Main execution
+log_info "Exporting data for autocompletes from the database."
+export_unique_values "SELECT DISTINCT industry_sector FROM public.companies WHERE industry_sector IS NOT NULL" "./data/export_docker/industry_sector.csv"
+export_unique_values "SELECT DISTINCT city FROM public.companies WHERE city IS NOT NULL" "./data/export_docker/city.csv"
+export_unique_values "SELECT DISTINCT legal_form FROM public.companies WHERE legal_form IS NOT NULL" "./data/export_docker/legal_form.csv"
+export_unique_values "SELECT DISTINCT region FROM public.companies WHERE region IS NOT NULL" "./data/export_docker/region.csv"
+
+log_info "Loading data for autocompletes into the database."
+transfer_csv_to_database "city" "./data/export_docker/city.csv" "name" ","
+transfer_csv_to_database "industry_sector" "./data/export_docker/industry_sector.csv" "name" ","
+transfer_csv_to_database "legal_form" "./data/export_docker/legal_form.csv" "name" ","
+transfer_csv_to_database "region" "./data/export_docker/region.csv" "name" ","
+
+log_success "Autocomplete data export and loading completed successfully."

scripts/load-csv-to-database.sh

@@ -18,86 +18,25 @@ transfer_csv_to_database() {
     local csv_file_path="$2"
     local columns="$3"
     local delimiter="$4"
-    local base_name
-    base_name=$(basename "$csv_file_path")
-    local container_csv_file="/tmp/$base_name"
 
     if [ ! -f "$csv_file_path" ]; then
         log_error "The CSV file '$csv_file_path' doesn't exist."
         exit 1
     fi
 
-    log_info "Transferring the CSV file '$csv_file_path' to the PostgreSQL database."
+    log_info "Copying data from '$csv_file_path' to the database table '$table_name'..."
 
-    cp "$csv_file_path" "$container_csv_file"
-    psql -d postgres -c "COPY $table_name($columns) FROM '$container_csv_file' DELIMITER '$delimiter' CSV HEADER;"
+    psql -d postgres -c "\copy $table_name($columns) FROM '$csv_file_path' DELIMITER '$delimiter' CSV HEADER;"
 
-    log_success "Transfer of '$container_csv_file' to the database table '$table_name' completed successfully."
+    log_success "Transfer of '$csv_file_path' to the database table '$table_name' completed successfully."
 }
 
-
-# Function to export unique values to a CSV or SQL file
-export_unique_values() {
-    local query="$1"
-    local output_file="$2"
-    local format="${3:-csv}" # Default format is CSV
-    local base_name
-    base_name=$(basename "$output_file")
-
-    mkdir -p "$(dirname "$output_file")"
-
-    if [ "$format" == "csv" ]; then
-        local output_csv="/tmp/$base_name.csv"
-        log_info "Exporting unique values for the $base_name table in CSV format..."
-
-        mkdir -p /tmp
-        psql -d postgres -c "\copy ($query) TO '$output_csv' CSV HEADER;"
-        cp "$output_csv" "$output_file"
-        log_success "Exported CSV file saved to $output_file"
-
-    elif [ "$format" == "sql" ]; then
-        log_info "Exporting unique values for the $base_name table in SQL format..."
-
-        local temp_table="temp_export"
-        local output_sql="$output_file.sql"
-
-        psql -d postgres -c "
-            DROP TABLE IF EXISTS $temp_table;
-            CREATE TABLE $temp_table AS
-            SELECT row_number() OVER () AS id, * FROM ($query) AS subquery;
-        "
-
-        pg_dump -U postgres --data-only --table="$temp_table" postgres >"$output_sql"
-        psql -d postgres -c "DROP TABLE IF EXISTS $temp_table;"
-
-        log_success "Exported SQL file saved to $output_sql"
-    else
-        log_error "Invalid format specified. Use 'csv' or 'sql'."
-        exit 1
-    fi
-}
-
-
 # Main script execution starts here
 log_info "Loading 'companies' and 'leaders' CSV data into the database."
-transfer_csv_to_database "companies" "./final.csv" "$(head -1 "./final.csv" | tr ';' ',')" ";"
-transfer_csv_to_database "leaders" "./leaders.csv" "$(head -1 "./leaders.csv" | tr ';' ',')" ";"
-
-log_info "Inserting scrapped data into the database."
-python3 "./scripts/load_scrapped_companies.py"
-
-log_info "Exporting data for autocompletes from the database."
-export_unique_values "SELECT DISTINCT industry_sector FROM public.companies" "./data/export_docker/industry_sector.csv"
-export_unique_values "SELECT DISTINCT city FROM public.companies" "./data/export_docker/city.csv"
-export_unique_values "SELECT DISTINCT legal_form FROM public.companies" "./data/export_docker/legal_form.csv"
-export_unique_values "SELECT DISTINCT region FROM public.companies" "./data/export_docker/region.csv"
-
-
-log_info "Loading data for autocompletes into the database."
-transfer_csv_to_database "city" "./data/export_docker/city.csv" "name" ","
-transfer_csv_to_database "industry_sector" "./data/export_docker/industry_sector.csv" "name" ","
-transfer_csv_to_database "legal_form" "./data/export_docker/legal_form.csv" "name" ","
-transfer_csv_to_database "region" "./data/export_docker/region.csv" "name" ","
+companies_header=$(head -n 1 "./final.csv" | tr ';' ',')
+leaders_header=$(head -n 1 "./leaders.csv" | tr ';' ',')
 
+transfer_csv_to_database "companies" "./final.csv" "$companies_header" ";"
+transfer_csv_to_database "leaders" "./leaders.csv" "$leaders_header" ";"
 
-log_success "Data loading and export completed successfully."
+log_success "CSV data loading completed successfully."

scripts/setup-db.sh

@@ -16,7 +16,17 @@
 set -e
 
 # Start Postgres in background
-docker-entrypoint.sh postgres &
+docker-entrypoint.sh postgres \
+  -c max_wal_size=6GB \
+  -c min_wal_size=1GB \
+  -c checkpoint_timeout=30min \
+  -c checkpoint_completion_target=0.9 \
+  -c synchronous_commit=off \
+  -c fsync=off \
+  -c full_page_writes=off \
+  -c wal_level=minimal \
+  -c max_wal_senders=0 \
+  &
 
 # Wait for Postgres to be ready
 echo "Waiting for Postgres..."
@@ -31,39 +41,50 @@ export DATABASE_URL="postgresql://postgres:root@localhost:5432/postgres"
 # Test connection
 psql -d postgres -c "SELECT 1;"
 
-# Run migrations
+# Run migrations (creates all tables + indexes)
 cd /app/schema
 pnpm exec prisma migrate deploy
 pnpm exec prisma generate
 cd /app
 
-psql -v ON_ERROR_STOP=1 --username="$POSTGRES_USER" <<EOSQL
-        ALTER SYSTEM SET max_wal_size = '6GB';
-        ALTER SYSTEM SET checkpoint_timeout = '30min';
-        ALTER SYSTEM SET synchronous_commit = off;
-        ALTER SYSTEM SET fsync = off;
-        ALTER SYSTEM SET full_page_writes = off;
-EOSQL
+# In order to speed up the CSV data load, we drop all non-PK indexes before loading,
+# then recreate them after. This is much faster than maintaining the indexes during the load.
+echo "Exporting current indexes to backup file..."
+psql -d postgres -f /app/scripts/sql/export-indexes.sql
 
-# Reload Postgres configuration
-psql -d postgres -c "SELECT pg_reload_conf();"
+echo "Dropping all non-PK indexes to optimize bulk data loading..."
+psql -d postgres -f /app/scripts/sql/drop-indexes-from-backup.sql
 
-# Load CSV
+echo "Loading CSV data into database (without indexes for maximum speed)..."
 /app/scripts/load-csv-to-database.sh
 
-# Reset performance settings
-psql -d postgres -v ON_ERROR_STOP=1 <<EOSQL
-        ALTER SYSTEM RESET max_wal_size;
-        ALTER SYSTEM RESET checkpoint_timeout;
-        ALTER SYSTEM RESET synchronous_commit;
-        ALTER SYSTEM RESET fsync;
-        ALTER SYSTEM RESET full_page_writes;
-EOSQL
+echo "Analyzing tables after CSV load..."
+psql -d postgres -c "ANALYZE public.companies;"
+psql -d postgres -c "ANALYZE public.leaders;"
+
+echo "Creating temporary index on siren_number for UPDATE optimization..."
+psql -d postgres -f /app/scripts/sql/create-temp-siren-index.sql
+
+echo "Updating companies with scraped data..."
+python3 "./scripts/load_scrapped_companies.py"
+
+echo "Dropping temporary siren_number index..."
+psql -d postgres -f /app/scripts/sql/drop-temp-siren-index.sql
+
+echo "Recreating all indexes from backup file..."
+psql -d postgres -f /app/scripts/sql/recreate-indexes-from-backup.sql
+
+echo "Verifying indexes match original state..."
+psql -d postgres -f /app/scripts/sql/verify-indexes.sql
+
+echo "Exporting DISTINCT values for autocomplete tables..."
+/app/scripts/export-and-load-autocomplete.sh
 
 # Reload Postgres configuration
 psql -d postgres -c "SELECT pg_reload_conf();"
 
-# Vacuum and analyze
+# Final VACUUM and analyze
+echo "Running final VACUUM and ANALYZE..."
 psql -d postgres -c "VACUUM FULL;"
 psql -d postgres -c "ANALYZE;"
 

scripts/sql/create-temp-siren-index.sql

@@ -0,0 +1,9 @@
+-- Create a temporary index on companies.siren_number to optimize the UPDATE operation
+-- in load_scrapped_companies.py which joins on siren_number
+-- This index is critical for UPDATE performance (converts O(n²) to O(n log n))
+
+CREATE INDEX IF NOT EXISTS ix_companies_siren_number_temp
+ON public.companies(siren_number);
+
+-- Analyze the table after creating the index to update statistics
+ANALYZE public.companies;

scripts/sql/drop-indexes-from-backup.sql

@@ -0,0 +1,45 @@
+-- Drop all indexes that were exported to the backup file
+-- This script drops only regular indexes, not those backing constraints
+
+DO $$
+DECLARE
+    index_record RECORD;
+    drop_count INTEGER := 0;
+    skip_count INTEGER := 0;
+BEGIN
+    -- Drop each index that's not a primary key
+    FOR index_record IN
+        SELECT i.indexname, i.tablename
+        FROM pg_indexes i
+        WHERE i.schemaname = 'public'
+          AND i.indexname NOT LIKE '%_pkey'
+        ORDER BY i.tablename, i.indexname
+    LOOP
+        BEGIN
+            -- Try to drop the index
+            EXECUTE format('DROP INDEX IF EXISTS public.%I', index_record.indexname);
+            drop_count := drop_count + 1;
+            RAISE NOTICE 'Dropped index: % on table %', index_record.indexname, index_record.tablename;
+        EXCEPTION
+            WHEN dependent_objects_still_exist THEN
+                -- Skip indexes that have dependencies (unique constraints, foreign keys, etc.)
+                skip_count := skip_count + 1;
+                RAISE NOTICE 'Skipped index: % on table % (has dependencies)', index_record.indexname, index_record.tablename;
+        END;
+    END LOOP;
+
+    RAISE NOTICE 'Total indexes dropped: %', drop_count;
+    RAISE NOTICE 'Total indexes skipped: %', skip_count;
+END $$;
+
+-- Verify all droppable indexes are gone
+SELECT
+    CASE
+        WHEN COUNT(*) = 0 THEN 'SUCCESS: All droppable indexes dropped'
+        ELSE 'WARNING: ' || COUNT(*) || ' droppable indexes still exist'
+    END as verification_result
+FROM pg_indexes i
+WHERE i.schemaname = 'public'
+  AND i.indexname NOT LIKE '%_pkey'
+  AND i.indexname NOT LIKE '%_key'
+  AND i.indexdef NOT LIKE 'CREATE UNIQUE%';

scripts/sql/drop-temp-siren-index.sql

@@ -0,0 +1,4 @@
+-- Drop the temporary siren_number index after the UPDATE operation completes
+-- This index will be recreated as part of create-indexes.sql with the proper name
+
+DROP INDEX IF EXISTS public.ix_companies_siren_number_temp;

scripts/sql/export-indexes.sql

@@ -0,0 +1,15 @@
+-- Export all non-primary-key, non-unique indexes to a file for later recreation
+-- This ensures we can recreate the exact same indexes after bulk data loading
+
+-- Create machine-readable CSV for recreation
+\copy (SELECT i.indexname, i.tablename, i.indexdef FROM pg_indexes i WHERE i.schemaname = 'public' AND i.indexname NOT LIKE '%_pkey' AND i.indexname NOT LIKE '%_key' AND i.indexdef NOT LIKE 'CREATE UNIQUE%' ORDER BY i.tablename, i.indexname) TO '/tmp/indexes_backup.csv' WITH (FORMAT CSV, HEADER true);
+
+-- Print summary
+SELECT
+    COUNT(*) as total_indexes,
+    COUNT(DISTINCT i.tablename) as affected_tables
+FROM pg_indexes i
+WHERE i.schemaname = 'public'
+  AND i.indexname NOT LIKE '%_pkey'
+  AND i.indexname NOT LIKE '%_key'
+  AND i.indexdef NOT LIKE 'CREATE UNIQUE%';