apb_extra_utils.postgres_pckg.psql_alchemy

  1#  coding=utf-8
  2#
  3#  Author: Ernesto Arredondo Martinez (ernestone@gmail.com)
  4#  Created: 31/03/2019
  5#  Copyright (c)
  6from __future__ import annotations
  7import re
  8from threading import Lock
  9
 10from apb_extra_utils.utils_logging import get_base_logger
 11from sqlalchemy.orm import sessionmaker
 12from sqlalchemy.engine import create_engine
 13from sqlalchemy.engine.url import make_url, URL
 14from sqlalchemy import MetaData, Table, text, Engine, inspect
 15from collections import namedtuple
 16
 17TYPE_UNKNOWN = 'UNKNOWN'
 18PG_DRIVER = 'postgresql+psycopg2'
 19
 20# Intenta importar soporte para tipos geométricos PostGIS
 21try:
 22    import geoalchemy2  # noqa: F401
 23
 24    HAS_GEOALCHEMY2 = True
 25except ImportError:
 26    HAS_GEOALCHEMY2 = False
 27
 28
 29def url_pg_string_connection(user='', psw='', host='localhost', port=5432, db='postgres', url_conn_string=None,
 30                             schemas='') -> URL:
 31    """
 32    Retorna una URL pg_connection parseada.
 33    Se pueden informar los parámetros uno a uno o directamente con la cadena de conexión completa (conn_string).
 34    En este ultimo caso se ignoran los parámetros individuales excepto el 'schemas'  que se añade a la cadena de conexión si no estuviera
 35
 36    Args:
 37        user (str): usuario
 38        psw (str): password
 39        host (str): host
 40        port (int): port
 41        db (str): database
 42        url_conn_string (str | URL): connection string or sqlalchemy URL
 43        schemas (str): schemas separados por comas
 44
 45    Returns:
 46        URL
 47    """
 48    if url_conn_string:
 49        url = make_url(url_conn_string)
 50    else:
 51        url = URL.create(
 52            drivername=PG_DRIVER,
 53            username=user or '', password=psw or '',
 54            host=host or '', port=port, database=db or '',
 55        )
 56
 57    if schemas:
 58        normalized_schemas = ','.join([s.strip() for s in str(schemas).split(',') if s.strip()])
 59        if normalized_schemas:
 60            url = url.update_query_dict(dict(options=f"-csearch_path={normalized_schemas}"))
 61
 62    return url
 63
 64
 65class EngPsqlAlchemy(object):
 66    """
 67    Clase que gestiona conexion sqlalchemy a Postgres
 68    """
 69    __slots__ = 'url_con_db', 'eng_db', 'logger', 'session_db', 'schemas'
 70    _CACHE = {}
 71    _CACHE_LOCK = Lock()
 72
 73    @classmethod
 74    def get_cached(cls, user=None, psw=None,
 75                   srvr_db='localhost', port_db=5432,
 76                   db='postgres', schemas=None,
 77                   a_logger=None, url_conn_string=None, force_new=False):
 78        """
 79        Recupera una instancia cacheada por conexión o crea una nueva si no existe.
 80
 81        La clave de caché se basa en user/host/port/db/schemas (o conn_string).
 82        Antes de devolver una instancia cacheada, valida y recupera su session si es necesario.
 83        """
 84        cache_key = url_pg_string_connection(user=user, psw=psw, host=srvr_db, port=port_db, db=db,
 85                                             url_conn_string=url_conn_string, schemas=schemas)
 86
 87        with cls._CACHE_LOCK:
 88            if force_new:
 89                cls._CACHE.pop(cache_key, None)
 90
 91            inst = cls._CACHE.get(cache_key)
 92            if inst is None:
 93                inst = cls(user=user, psw=psw, srvr_db=srvr_db, port_db=port_db, db=db,
 94                           schemas=schemas, a_logger=a_logger, url_conn_string=url_conn_string)
 95                cls._CACHE[cache_key] = inst
 96            else:
 97                inst.ensure_session()
 98
 99            return inst
100
101    def __init__(self, user=None, psw=None, srvr_db='localhost', port_db=5432, db='postgres', schemas=None,
102                 a_logger=None, url_conn_string=None):
103        """
104        Retorna engine para database postgres. Si no se informa ningun argumento retorna el cacheado
105
106        Args:
107            user (str=None):
108            psw (str=None):
109            srvr_db (str=None):
110            port_db (int=None):
111            db (str=None):
112            schemas (str=None): schemas separated by comma
113            a_logger (logging.Logger=None):
114            url_conn_string (str|URL=None): connection string or sqlalchemy URL
115
116        Returns:
117            sqlalchemy.engine.base.Engine
118        """
119        if (user is None or psw is None) and url_conn_string is None:
120            raise ValueError("Debe informar user y psw, o alternativamente conn_string")
121
122        self.url_con_db = url_pg_string_connection(user=user, psw=psw, host=srvr_db, port=port_db, db=db,
123                                                   url_conn_string=url_conn_string, schemas=schemas)
124
125        if schemas:
126            self.schemas = [s.strip() for s in str(schemas).split(',') if s.strip()]
127
128        self.logger = a_logger
129        self.__set_logger()
130
131        self.logger.debug(f"Initializing EngPsqlAlchemy '{self.url_con_db}'")
132
133        extra_args = {}
134        if HAS_GEOALCHEMY2:
135            extra_args['plugins'] = ['geoalchemy2']
136
137        eng_db: Engine = create_engine(
138            self.url_con_db,
139            **extra_args
140        )
141
142        eng_db.connect()
143
144        self.eng_db = eng_db
145        # self.eng_db.logger = self.logger
146        self._set_session()
147
148    @classmethod
149    def from_url_conn_string(cls, url_conn_string, a_logger=None):
150        """
151        Constructor alternativo a partir de cadena de conexion.
152
153        Args:
154            url_conn_string (str | URL):
155            a_logger (logging.Logger=None):
156
157        Returns:
158            EngPsqlAlchemy
159        """
160        return cls(url_conn_string=url_conn_string, a_logger=a_logger)
161
162    def __set_logger(self):
163        """
164        Asigna el LOGGER po defecto si este no se ha informado al inicializar el gestor
165
166        Returns:
167        """
168        if self.logger is None:
169            self.logger = get_base_logger(f'{self.__class__.__name__}({self.url_con_db})')
170
171    def _set_session(self, eng_db=None):
172        """
173        Configura session para controlar workflow de transacciones (commit, rollback, close...)
174
175        Args:
176            eng_db:
177        """
178        session = sessionmaker()
179        session.configure(bind=eng_db if eng_db else self.eng_db)
180        self.session_db = session()
181
182    def is_session_alive(self):
183        """Valida si la session actual responde correctamente con un ping simple."""
184        if not getattr(self, 'session_db', None):
185            return False
186
187        try:
188            self.session_db.execute(text('SELECT 1'))
189            return True
190        except Exception as exc:
191            if self.logger:
192                self.logger.warning(f"Session no disponible, se intentará recuperar: {exc}")
193            return False
194
195    def recover_session(self, dispose_engine=False):
196        """Recrea la sesión, opcionalmente reciclando el pool del engine."""
197        try:
198            if getattr(self, 'session_db', None):
199                self.session_db.close()
200        except Exception:
201            pass
202
203        if dispose_engine and getattr(self, 'eng_db', None):
204            self.eng_db.dispose()
205
206        self._set_session()
207        return self.session_db
208
209    def ensure_session(self):
210        """Garantiza que la session está operativa; si no, la recupera."""
211        if self.is_session_alive():
212            return self.session_db
213
214        self.recover_session(dispose_engine=False)
215        if self.is_session_alive():
216            return self.session_db
217
218        # Último intento reciclando conexiones del pool.
219        self.recover_session(dispose_engine=True)
220        return self.session_db
221
222    def __del__(self):
223        """
224        Cierra la conexion al matar la instancia
225        """
226        try:
227            if hasattr(self, 'session_db'):
228                self.session_db.close()
229            if hasattr(self, "con_db"):
230                self.eng_db = None
231        except:
232            pass
233
234    def commit(self):
235        """
236        Hace commit sobre la sesion actual
237        """
238        if self.session_db:
239            self.session_db.commit()
240
241    def rollback(self):
242        """
243        Hace rollback sobre la sesion actual
244        """
245        if self.session_db:
246            self.session_db.rollback()
247
248    def iter_rows_result(self, query_res, nom_row=None):
249        """
250        Itera las filas del resultado como namedtuple con las claves de los campos seleccionados
251
252        Args:
253            query_res (sqlalchemy.engine.result.ResultProxy)
254            nom_row (str=None): nombre de la clase namedtuple si aplica
255
256        Yields:
257            namedtuple o tuple
258        """
259        row_dd = None
260        if query_res.keys():
261            row_dd = namedtuple(nom_row if nom_row else 'row_result',
262                                [n_col.replace(" ", "_") for n_col in query_res.keys()])
263
264        for row in query_res:
265            yield row_dd(*row) if row_dd else row
266
267    def table_view(self, nom_tab_view, schema=None):
268        """
269        Retorna acceso a tabla sobre engine DB (
270
271        Args:
272            nom_tab_view (str):
273            schema (str=None):
274
275        Returns:
276
277        """
278        schema = self.check_exists_table_view(nom_tab_view, schema=schema)
279
280        meta = MetaData()
281
282        extra_args = {}
283        if schema:
284            extra_args['schema'] = schema
285
286        a_tab = Table(nom_tab_view, meta, autoload_with=self.eng_db, **extra_args)
287
288        return a_tab
289
290    def tables_views(self, schema=None, only_tables=False) -> dict[str, list[str]]:
291        """
292        Lista las tablas/views disponibles para la conexion actual.
293
294        Args:
295            only_tables:
296            schema (str=None): esquema a consultar; si es None, usa el search_path/por defecto.
297
298        Returns:
299            dict[str, list[str]]: diccionario con las tablas/views disponibles por esquema
300        """
301        inspector = inspect(self.eng_db)
302        if not schema:
303            schemas = inspector.get_schema_names()
304        else:
305            schemas = [schema]
306
307        tables_views = {}
308        for schema in schemas:
309            tabs_schema = sorted(inspector.get_table_names(schema=schema))
310            if not only_tables:
311                tabs_schema.extend(sorted(inspector.get_view_names(schema=schema)))
312            if not tabs_schema:
313                continue
314            tables_views[schema] = tabs_schema
315
316        return tables_views
317
318    def columns_table_view(self, nom_tab, schema=None):
319        """
320        Retorna columnas y tipos de una tabla o vista.
321
322        Args:
323            nom_tab (str): nombre de tabla o vista.
324            schema (str=None): esquema a consultar.
325
326        Returns:
327            dict[str, str]: diccionario indexado por nombre de columna con el tipo.
328        """
329        schema = self.check_exists_table_view(nom_tab, schema=schema)
330
331        # noinspection SqlResolve,SqlNoDataSourceInspection
332        # language=PostgreSQL
333        sql_query = text("""
334                         SELECT a.attname                                       AS column_name,
335                                pg_catalog.format_type(a.atttypid, a.atttypmod) AS column_type
336                         FROM pg_catalog.pg_attribute a
337                                  INNER JOIN pg_catalog.pg_class c
338                                             ON c.oid = a.attrelid
339                                  INNER JOIN pg_catalog.pg_namespace n
340                                             ON n.oid = c.relnamespace
341                         WHERE a.attnum > 0
342                           AND NOT a.attisdropped
343                           AND c.relkind IN ('r', 'p', 'v', 'm', 'f')
344                           AND c.relname = :table_name
345                           AND (
346                             (:schema IS NOT NULL AND n.nspname = :schema)
347                                 OR (:schema IS NULL AND pg_catalog.pg_table_is_visible(c.oid))
348                             )
349                         ORDER BY a.attnum
350                         """)
351
352        res = self.ensure_session().execute(sql_query, {
353            'table_name': nom_tab,
354            'schema': schema,
355        })
356
357        cols = {}
358        for row in res:
359            cols[row.column_name] = row.column_type.replace('public.', '') if row.column_type else TYPE_UNKNOWN
360
361        return cols
362
363    def check_exists_table_view(self, nom_tab_or_view, schema=None) -> str:
364        """
365        Valida que exista una tabla o vista para la conexión actual y devuelve su esquema
366
367        Args:
368            nom_tab_or_view (str): nombre de tabla o vista.
369            schema (str=None): esquema a consultar.
370
371        Returns:
372            str: schema de la tabla o vista.
373        """
374        schema_tab_view = None
375        for sch, tabs_views in self.tables_views(schema=schema).items():
376            if nom_tab_or_view in tabs_views:
377                schema_tab_view = sch
378                break
379
380        if not schema_tab_view:
381            raise ValueError(f"No existe la tabla/vista '{nom_tab_or_view}' "
382                             f"{'para el esquema {} '.format(schema) if schema else ''}"
383                             f"sobre la conexión '{self.url_con_db}'")
384
385        return schema_tab_view
386
387    def geoms_table_view(self, nom_tab, schema=None):
388        """
389        Retorna columnas geométricas/geográficas de una tabla o vista con metadatos detallados.
390
391        Args:
392            nom_tab (str): nombre de tabla o vista.
393            schema (str=None): esquema a consultar
394
395        Returns:
396            dict[str, dict]: diccionario indexado por nombre de columna con metadatos geométricos.
397        """
398        schema = self.check_exists_table_view(nom_tab, schema=schema)
399
400        # noinspection SqlResolve,SqlNoDataSourceInspection
401        # language=PostgreSQL
402        sql_query = text("""
403                         SELECT a.attname                                                               AS column_name,
404                                REPLACE(pg_catalog.format_type(a.atttypid, a.atttypmod), 'public.', '') AS column_type,
405                                t.typname                                                               AS base_type
406                         FROM pg_catalog.pg_attribute a
407                                  INNER JOIN pg_catalog.pg_class c
408                                             ON c.oid = a.attrelid
409                                  INNER JOIN pg_catalog.pg_namespace n
410                                             ON n.oid = c.relnamespace
411                                  INNER JOIN pg_catalog.pg_type t
412                                             ON t.oid = a.atttypid
413                         WHERE a.attnum > 0
414                           AND NOT a.attisdropped
415                           AND c.relkind IN ('r', 'p', 'v', 'm', 'f')
416                           AND c.relname = :table_name
417                           AND t.typname IN ('geometry', 'geography')
418                           AND (
419                             (:schema IS NOT NULL AND n.nspname = :schema)
420                                 OR (:schema IS NULL AND pg_catalog.pg_table_is_visible(c.oid))
421                             )
422                         ORDER BY a.attnum
423                         """)
424
425        res = self.ensure_session().execute(sql_query, {
426            'table_name': nom_tab,
427            'schema': schema,
428        })
429
430        geoms = {}
431        for row in res:
432            type_str, base_type, geometry_type, srid = self.__parse_spatial_type(
433                row.column_type,
434                base_type_hint=row.base_type,
435            )
436            geoms[row.column_name] = {
437                'type': type_str,
438                'base_type': base_type,
439                'geometry_type': geometry_type,
440                'srid': srid,
441            }
442
443        return geoms
444
445    @staticmethod
446    def __parse_spatial_type(type_str, base_type_hint=None):
447        """Parsea tipos espaciales tipo geometry(Point,25831) sin depender de funciones PostGIS."""
448        if not type_str and not base_type_hint:
449            return TYPE_UNKNOWN, TYPE_UNKNOWN, TYPE_UNKNOWN, None
450
451        clean_type = (type_str or '').replace('public.', '').strip() or TYPE_UNKNOWN
452        base_type = (base_type_hint or clean_type.split('(', 1)[0]).strip().lower()
453        geometry_type = TYPE_UNKNOWN
454        srid = None
455
456        # Ejemplos esperados: geometry(Point,25831), geography(MultiPolygon,4326), geometry
457        m = re.match(r'^(geometry|geography)\(([^,()]+)(?:\s*,\s*(\d+))?\)$', clean_type, flags=re.IGNORECASE)
458        if m:
459            base_type = m.group(1).lower()
460            geometry_type = m.group(2)
461            srid = int(m.group(3)) if m.group(3) else None
462        elif clean_type.lower() in ('geometry', 'geography'):
463            base_type = clean_type.lower()
464
465        if base_type not in ('geometry', 'geography'):
466            base_type = TYPE_UNKNOWN
467
468        return clean_type, base_type, geometry_type, srid
469
470    def rows_table_view(self, nom_tab, sql_query=None, **table_args):
471        """
472        Itera sobre los registros de la tabla
473
474        Args:
475            nom_tab (str):
476            sql_query (str=None):
477            **table_args: argumentos de la funcion table()
478
479        Yields:
480            namedtuple
481        """
482        tab = self.table_view(nom_tab, **table_args)
483
484        query = tab.select()
485        if sql_query:
486            query = query.where(text(sql_query))
487
488        res = self.ensure_session().execute(query)
489
490        for row in self.iter_rows_result(res, tab.name):
491            yield row
492
493    def insert_rows_table(self, nom_tab, row_values, **table_args):
494        """
495        Inserta los registros en la tabla
496
497        Args:
498            nom_tab (str):
499            row_values (list): lista de dicts con los valores de los campos a insertar por cada fila
500            **table_args: argumentos de la funcion table()
501
502        Returns:
503            list: lista de los rows_inserted (namedtuple)
504        """
505        tab = self.table_view(nom_tab, **table_args)
506        query = tab.insert(values=row_values).returning(*tab.columns())
507
508        res = self.ensure_session().execute(query)
509
510        rows_inserted = [*self.iter_rows_result(res, tab.name)]
511
512        return rows_inserted
513
514    def update_rows_table(self, tab, values, sql_query=None, **table_args):
515        """
516        Actualiza los registros de la tabla que cumplan con where_query (todos por defecto)
517
518        Args:
519            tab (str):
520            values (dict):
521            sql_query (str):
522            **table_args: argumentos de la funcion table()
523
524        Returns:
525            generator: iterador con la lista de elementos a devolver
526        """
527        tab = self.table_view(tab, **table_args)
528        query = tab.update()
529        if sql_query:
530            query = query.where(text(sql_query))
531        query = query.values(values).returning(*tab.columns())
532
533        res = self.ensure_session().execute(query)
534
535        return res
536
537    def remove_rows_table(self, nom_tab, sql_query=None, **table_args):
538        """
539        Borra los registros de la tabla que cumplan con where_query (todos por defecto)
540
541        Args:
542            nom_tab (str):
543            sql_query (str):
544            **table_args: argumentos de la funcion table()
545
546        Returns:
547            res
548        """
549        tab = self.table_view(nom_tab, **table_args)
550        query = tab.delete()
551
552        if sql_query:
553            query = query.where(text(sql_query))
554
555        res = self.ensure_session().execute(query)
556
557        return res
558
559
560if __name__ == '__main__':
561    from fire import Fire
562
563    Fire()