<commit_before>from alembic import op, context
import sqlalchemy as sa
import logging
from contextlib import ExitStack
from tempfile import TemporaryDirectory
from shutil import copyfileobj
from urllib.request import urlopen
import lzma
import os, os.path


schema = 'ref_geo'

logger = logging.getLogger('alembic.runtime.migration')


class open_geofile(ExitStack):
    def __init__(self, base_url, filename):
        super().__init__()
        self.base_url = base_url
        self.filename = filename
        self.geo_dir = context.get_x_argument(as_dictionary=True).get('geo-data-directory')

    def __enter__(self):
        stack = super().__enter__()
        if not self.geo_dir:
            self.geo_dir = stack.enter_context(TemporaryDirectory())
            logger.info("Created temporary directory '{}'".format(self.geo_dir))
        if not os.path.exists(self.geo_dir):
            os.mkdir(self.geo_dir)
        geofile_path = os.path.join(self.geo_dir, self.filename)
        if not os.path.isfile(geofile_path):
            logger.info("Downloading '{}'…".format(self.filename))
            with urlopen('{}{}'.format(self.base_url, self.filename)) as response, \
                                              open(geofile_path, 'wb') as geofile:
                copyfileobj(response, geofile)
        return stack.enter_context(lzma.open(geofile_path))


def delete_area_with_type(area_type):
    op.execute(f"""
        DELETE FROM {schema}.l_areas la
        USING {schema}.get_id_area_type('{area_type}') as area_type
        WHERE la.id_type = area_type
    """)


def create_temporary_grids_table(schema, temp_table_name):
    logger.info("Create temporary grids table…")
    op.execute(f"""
        CREATE TABLE {schema}.{temp_table_name} (
            gid integer NOT NULL,
            cd_sig character varying(21),
            code character varying(10),
            geom public.geometry(MultiPolygon,2154),
            geojson character varying
        )
    """)
    op.execute(f"""
        ALTER TABLE ONLY {schema}.{temp_table_name}
            ADD CONSTRAINT {temp_table_name}_pkey PRIMARY KEY (gid)
    """)


def insert_grids_and_drop_temporary_table(schema, temp_table_name, area_type):
    logger.info("Copy grids in l_areas…")
    op.execute(f"""<commit_msg>add comments on ref geo utils<commit_after>from alembic import op, context
import sqlalchemy as sa
import logging
from contextlib import ExitStack
from tempfile import TemporaryDirectory
from shutil import copyfileobj
from urllib.request import urlopen
import lzma
import os, os.path


schema = 'ref_geo'

logger = logging.getLogger('alembic.runtime.migration')


"""
Ce contextmanager permet d’ouvrir un fichier compressé (xz) après l’avoir préalablement téléchargé.
Le fichier téléchargé est enregistré dans le dossier spécifié par -x geo-data-directory=…
Si aucun dossier n’est spécifié, un dossier temporaire, supprimé à la fin de la migration, est utilisé.
"""
class open_geofile(ExitStack):
    def __init__(self, base_url, filename):
        super().__init__()
        self.base_url = base_url
        self.filename = filename
        self.geo_dir = context.get_x_argument(as_dictionary=True).get('geo-data-directory')

    def __enter__(self):
        stack = super().__enter__()
        if not self.geo_dir:
            self.geo_dir = stack.enter_context(TemporaryDirectory())
            logger.info("Created temporary directory '{}'".format(self.geo_dir))
        if not os.path.exists(self.geo_dir):
            os.mkdir(self.geo_dir)
        geofile_path = os.path.join(self.geo_dir, self.filename)
        if not os.path.isfile(geofile_path):
            logger.info("Downloading '{}'…".format(self.filename))
            with urlopen('{}{}'.format(self.base_url, self.filename)) as response, \
                                              open(geofile_path, 'wb') as geofile:
                copyfileobj(response, geofile)
        return stack.enter_context(lzma.open(geofile_path))


"""
Supprimer les zones d’un type donnée, e.g. 'DEP', 'COM', …
"""
def delete_area_with_type(area_type):
    op.execute(f"""
        DELETE FROM {schema}.l_areas la
        USING {schema}.get_id_area_type('{area_type}') as area_type
        WHERE la.id_type = area_type
    """)


def create_temporary_grids_table(schema, temp_table_name):
    logger.info("Create temporary grids table…")
    op.execute(f"""
        CREATE TABLE {schema}.{temp_table_name} (
            gid integer NOT NULL,
            cd_sig character varying(21),
            code character varying(10),
            geom public.geometry(MultiPolygon,2154),
            geojson character varying
        )
    """)
    op.execute(f"""
        ALTER TABLE ONLY {schema}.{temp_table_name}
            ADD CONSTRAINT {temp_table_name}_pkey PRIMARY KEY (gid)
    """)


def insert_grids_and_drop_temporary_table(schema, temp_table_name, area_type):
    logger.info("Copy grids in l_areas…")
    op.execute(f"""
{'fraction_non_alphanumeric': 0.10625880458844839, 'fraction_numerical': 0.0032199637754075265, 'mean_word_length': 2.976, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4617433', 'n_tokens_mistral': 1588, 'n_tokens_neox': 1534, 'n_words': 349}