meerschaum 3.0.0rc1__py3-none-any.whl → 3.0.0rc3__py3-none-any.whl

meerschaum/connectors/sql/_pipes.py

@@ -25,7 +25,6 @@ def register_pipe(
     Register a new pipe.
     A pipe's attributes must be set before registering.
     """
-    from meerschaum.utils.debug import dprint
     from meerschaum.utils.packages import attempt_import
     from meerschaum.utils.sql import json_flavors
 
@@ -148,7 +147,7 @@ def fetch_pipes_keys(
     tags: Optional[List[str]] = None,
     params: Optional[Dict[str, Any]] = None,
     debug: bool = False
-) -> Optional[List[Tuple[str, str, Optional[str]]]]:
+) -> List[Tuple[str, str, Optional[str]]]:
     """
     Return a list of tuples corresponding to the parameters provided.
 
@@ -163,17 +162,27 @@ def fetch_pipes_keys(
     location_keys: Optional[List[str]], default None
         List of location_keys to search by.
 
+    tags: Optional[List[str]], default None
+        List of pipes to search by.
+
     params: Optional[Dict[str, Any]], default None
         Dictionary of additional parameters to search by.
         E.g. `--params pipe_id:1`
 
     debug: bool, default False
         Verbosity toggle.
+
+    Returns
+    -------
+    A list of tuples of pipes' keys (connector_keys, metric_key, location_key).
     """
-    from meerschaum.utils.debug import dprint
     from meerschaum.utils.packages import attempt_import
     from meerschaum.utils.misc import separate_negation_values
-    from meerschaum.utils.sql import OMIT_NULLSFIRST_FLAVORS, table_exists
+    from meerschaum.utils.sql import (
+        OMIT_NULLSFIRST_FLAVORS,
+        table_exists,
+        json_flavors,
+    )
     from meerschaum._internal.static import STATIC_CONFIG
     import json
     from copy import deepcopy
@@ -261,25 +270,49 @@ def fetch_pipes_keys(
     in_ex_tag_groups = [separate_negation_values(tag_group) for tag_group in tag_groups]
 
     ors, nands = [], []
-    for _in_tags, _ex_tags in in_ex_tag_groups:
-        sub_ands = []
-        for nt in _in_tags:
-            sub_ands.append(
-                sqlalchemy.cast(
-                    pipes_tbl.c['parameters'],
-                    sqlalchemy.String,
-                ).like(f'%"tags":%"{nt}"%')
-            )
-        if sub_ands:
-            ors.append(sqlalchemy.and_(*sub_ands))
-
-        for xt in _ex_tags:
-            nands.append(
-                sqlalchemy.cast(
-                    pipes_tbl.c['parameters'],
-                    sqlalchemy.String,
-                ).not_like(f'%"tags":%"{xt}"%')
-            )
+    if self.flavor in json_flavors:
+        from sqlalchemy.dialects import postgresql
+        for _in_tags, _ex_tags in in_ex_tag_groups:
+            if _in_tags:
+                ors.append(
+                    sqlalchemy.and_(
+                        pipes_tbl.c['parameters'].cast(postgresql.JSONB).has_key('tags'),
+                        pipes_tbl.c['parameters']['tags'].cast(
+                            postgresql.JSONB
+                        ).contains(_in_tags)
+                    )
+                )
+            for xt in _ex_tags:
+                nands.append(
+                    sqlalchemy.not_(
+                        sqlalchemy.and_(
+                            pipes_tbl.c['parameters'].cast(postgresql.JSONB).has_key('tags'),
+                            pipes_tbl.c['parameters']['tags'].cast(
+                                postgresql.JSONB
+                            ).contains([xt])
+                        )
+                    )
+                )
+    else:
+        for _in_tags, _ex_tags in in_ex_tag_groups:
+            sub_ands = []
+            for nt in _in_tags:
+                sub_ands.append(
+                    sqlalchemy.cast(
+                        pipes_tbl.c['parameters'],
+                        sqlalchemy.String,
+                    ).like(f'%"tags":%"{nt}"%')
+                )
+            if sub_ands:
+                ors.append(sqlalchemy.and_(*sub_ands))
+
+            for xt in _ex_tags:
+                nands.append(
+                    sqlalchemy.cast(
+                        pipes_tbl.c['parameters'],
+                        sqlalchemy.String,
+                    ).not_like(f'%"tags":%"{xt}"%')
+                )
 
     q = q.where(sqlalchemy.and_(*nands)) if nands else q
     q = q.where(sqlalchemy.or_(*ors)) if ors else q
@@ -294,7 +327,7 @@ def fetch_pipes_keys(
 
     ### execute the query and return a list of tuples
     if debug:
-        dprint(q.compile(compile_kwargs={'literal_binds': True}))
+        dprint(q)
     try:
         rows = (
             self.execute(q).fetchall()
@@ -338,7 +371,6 @@ def create_indices(
     """
     Create a pipe's indices.
     """
-    from meerschaum.utils.debug import dprint
     if debug:
         dprint(f"Creating indices for {pipe}...")
 
@@ -392,7 +424,6 @@ def drop_indices(
     """
     Drop a pipe's indices.
     """
-    from meerschaum.utils.debug import dprint
     if debug:
         dprint(f"Dropping indices for {pipe}...")
 
@@ -1008,6 +1039,8 @@ def get_pipe_data(
     limit: Optional[int] = None,
     begin_add_minutes: int = 0,
     end_add_minutes: int = 0,
+    chunksize: Optional[int] = -1,
+    as_iterator: bool = False,
     debug: bool = False,
     **kw: Any
 ) -> Union[pd.DataFrame, None]:
@@ -1044,14 +1077,17 @@ def get_pipe_data(
         If specified, limit the number of rows retrieved to this value.
 
     begin_add_minutes: int, default 0
-        The number of minutes to add to the `begin` datetime (i.e. `DATEADD`.
+        The number of minutes to add to the `begin` datetime (i.e. `DATEADD`).
 
     end_add_minutes: int, default 0
-        The number of minutes to add to the `end` datetime (i.e. `DATEADD`.
+        The number of minutes to add to the `end` datetime (i.e. `DATEADD`).
 
     chunksize: Optional[int], default -1
         The size of dataframe chunks to load into memory.
 
+    as_iterator: bool, default False
+        If `True`, return the chunks iterator directly.
+
     debug: bool, default False
         Verbosity toggle.
 
@@ -1060,43 +1096,58 @@ def get_pipe_data(
     A `pd.DataFrame` of the pipe's data.
 
     """
-    import json
-    from meerschaum.utils.misc import parse_df_datetimes, to_pandas_dtype
+    import functools
     from meerschaum.utils.packages import import_pandas
-    from meerschaum.utils.dtypes import (
-        attempt_cast_to_numeric,
-        attempt_cast_to_uuid,
-        attempt_cast_to_bytes,
-        attempt_cast_to_geometry,
-        are_dtypes_equal,
-    )
+    from meerschaum.utils.dtypes import to_pandas_dtype, are_dtypes_equal
     from meerschaum.utils.dtypes.sql import get_pd_type_from_db_type
     pd = import_pandas()
     is_dask = 'dask' in pd.__name__
 
     cols_types = pipe.get_columns_types(debug=debug) if pipe.enforce else {}
+    pipe_dtypes = pipe.get_dtypes(infer=False, debug=debug) if pipe.enforce else {}
+
+    remote_pandas_types = {
+        col: to_pandas_dtype(get_pd_type_from_db_type(typ))
+        for col, typ in cols_types.items()
+    }
+    remote_dt_cols_types = {
+        col: typ
+        for col, typ in remote_pandas_types.items()
+        if are_dtypes_equal(typ, 'datetime')
+    }
+    remote_dt_tz_aware_cols_types = {
+        col: typ
+        for col, typ in remote_dt_cols_types.items()
+        if ',' in typ or typ == 'datetime'
+    }
+    remote_dt_tz_naive_cols_types = {
+        col: typ
+        for col, typ in remote_dt_cols_types.items()
+        if col not in remote_dt_tz_aware_cols_types
+    }
+
+    configured_pandas_types = {
+        col: to_pandas_dtype(typ)
+        for col, typ in pipe_dtypes.items()
+    }
+    configured_lower_precision_dt_cols_types = {
+        col: typ
+        for col, typ in pipe_dtypes.items()
+        if (
+            are_dtypes_equal('datetime', typ)
+            and '[' in typ
+            and 'ns' not in typ
+        )
+        
+    }
+
     dtypes = {
-        **{
-            col: get_pd_type_from_db_type(typ)
-            for col, typ in cols_types.items()
-        },
-        **{
-            p_col: to_pandas_dtype(p_typ)
-            for p_col, p_typ in pipe.dtypes.items()
-        },
+        **remote_pandas_types,
+        **configured_pandas_types,
+        **remote_dt_tz_aware_cols_types,
+        **remote_dt_tz_naive_cols_types,
+        **configured_lower_precision_dt_cols_types
     } if pipe.enforce else {}
-    if dtypes:
-        if self.flavor == 'sqlite':
-            if not pipe.columns.get('datetime', None):
-                _dt = pipe.guess_datetime()
-            else:
-                _dt = pipe.get_columns('datetime')
-
-            if _dt:
-                dt_type = dtypes.get(_dt, 'object').lower()
-                if 'datetime' not in dt_type:
-                    if 'int' not in dt_type:
-                        dtypes[_dt] = 'datetime64[ns, UTC]'
 
     existing_cols = cols_types.keys()
     select_columns = (
@@ -1113,13 +1164,20 @@ def get_pipe_data(
             and col not in (omit_columns or [])
         ]
     ) if pipe.enforce else select_columns
+
     if select_columns:
         dtypes = {col: typ for col, typ in dtypes.items() if col in select_columns}
+
     dtypes = {
-        col: to_pandas_dtype(typ)
+        col: typ
         for col, typ in dtypes.items()
-        if col in select_columns and col not in (omit_columns or [])
+        if col in (select_columns or [col]) and col not in (omit_columns or [])
     } if pipe.enforce else {}
+
+    if debug:
+        dprint(f"[{self}] `read()` dtypes:")
+        mrsm.pprint(dtypes)
+
     query = self.get_pipe_data_query(
         pipe,
         select_columns=select_columns,
@@ -1135,91 +1193,25 @@ def get_pipe_data(
         **kw
     )
 
+    read_kwargs = {}
     if is_dask:
         index_col = pipe.columns.get('datetime', None)
-        kw['index_col'] = index_col
+        read_kwargs['index_col'] = index_col
 
-    numeric_columns = [
-        col
-        for col, typ in pipe.dtypes.items()
-        if typ.startswith('numeric') and col in dtypes
-    ]
-    uuid_columns = [
-        col
-        for col, typ in pipe.dtypes.items()
-        if typ == 'uuid' and col in dtypes
-    ]
-    bytes_columns = [
-        col
-        for col, typ in pipe.dtypes.items()
-        if typ == 'bytes' and col in dtypes
-    ]
-    geometry_columns = [
-        col
-        for col, typ in pipe.dtypes.items()
-        if typ.startswith('geometry') and col in dtypes
-    ]
-
-    kw['coerce_float'] = kw.get('coerce_float', (len(numeric_columns) == 0))
-
-    df = self.read(
+    chunks = self.read(
         query,
+        chunksize=chunksize,
+        as_iterator=True,
+        coerce_float=False,
         dtype=dtypes,
         debug=debug,
-        **kw
+        **read_kwargs
     )
-    for col in numeric_columns:
-        if col not in df.columns:
-            continue
-        df[col] = df[col].apply(attempt_cast_to_numeric)
 
-    for col in uuid_columns:
-        if col not in df.columns:
-            continue
-        df[col] = df[col].apply(attempt_cast_to_uuid)
-
-    for col in bytes_columns:
-        if col not in df.columns:
-            continue
-        df[col] = df[col].apply(attempt_cast_to_bytes)
+    if as_iterator:
+        return chunks
 
-    for col in geometry_columns:
-        if col not in df.columns:
-            continue
-        df[col] = df[col].apply(attempt_cast_to_geometry)
-
-    if self.flavor == 'sqlite':
-        ignore_dt_cols = [
-            col
-            for col, dtype in pipe.dtypes.items()
-            if not are_dtypes_equal(str(dtype), 'datetime')
-        ]
-        ### NOTE: We have to consume the iterator here to ensure that datetimes are parsed correctly
-        df = (
-            parse_df_datetimes(
-                df,
-                ignore_cols=ignore_dt_cols,
-                chunksize=kw.get('chunksize', None),
-                strip_timezone=(pipe.tzinfo is None),
-                debug=debug,
-            ) if isinstance(df, pd.DataFrame) else (
-                [
-                    parse_df_datetimes(
-                        c,
-                        ignore_cols=ignore_dt_cols,
-                        chunksize=kw.get('chunksize', None),
-                        strip_timezone=(pipe.tzinfo is None),
-                        debug=debug,
-                    )
-                    for c in df
-                ]
-            )
-        )
-        for col, typ in dtypes.items():
-            if typ != 'json':
-                continue
-            df[col] = df[col].apply(lambda x: json.loads(x) if x is not None else x)
-    return df
+    return pd.concat(chunks)
 
 
 def get_pipe_data_query(
@@ -1552,13 +1544,7 @@ def create_pipe_table_from_df(
     """
     Create a pipe's table from its configured dtypes and an incoming dataframe.
     """
-    from meerschaum.utils.dataframe import (
-        get_json_cols,
-        get_numeric_cols,
-        get_uuid_cols,
-        get_datetime_cols,
-        get_bytes_cols,
-    )
+    from meerschaum.utils.dataframe import get_special_cols
     from meerschaum.utils.sql import (
         get_create_table_queries,
         sql_item_name,
@@ -1587,30 +1573,7 @@ def create_pipe_table_from_df(
             for col_ix, col in pipe.columns.items()
             if col and col_ix != 'primary'
         },
-        **{
-            col: 'uuid'
-            for col in get_uuid_cols(df)
-        },
-        **{
-            col: 'json'
-            for col in get_json_cols(df)
-        },
-        **{
-            col: 'numeric'
-            for col in get_numeric_cols(df)
-        },
-        **{
-            col: 'bytes'
-            for col in get_bytes_cols(df)
-        },
-        **{
-            col: 'datetime64[ns, UTC]'
-            for col in get_datetime_cols(df, timezone_aware=True, timezone_naive=False)
-        },
-        **{
-            col: 'datetime64[ns]'
-            for col in get_datetime_cols(df, timezone_aware=False, timezone_naive=True)
-        },
+        **get_special_cols(df),
         **pipe.dtypes
     }
     autoincrement = (
@@ -1762,18 +1725,16 @@ def sync_pipe(
             _ = pipe.__dict__.pop('_columns_types', None)
             if not self.exec_queries(alter_cols_queries, debug=debug):
                 warn(f"Failed to alter columns for {pipe}.")
-            else:
-                _ = pipe.infer_dtypes(persist=True)
 
     ### NOTE: Oracle SQL < 23c (2023) and SQLite does not support booleans,
     ### so infer bools and persist them to `dtypes`.
     if self.flavor in ('oracle', 'sqlite', 'mysql', 'mariadb'):
-        pipe_dtypes = pipe.dtypes
+        pipe_dtypes = pipe.get_dtypes(infer=False, debug=debug)
         new_bool_cols = {
             col: 'bool[pyarrow]'
             for col, typ in df.dtypes.items()
             if col not in pipe_dtypes
-            and are_dtypes_equal(str(typ), 'bool')
+                and are_dtypes_equal(str(typ), 'bool')
         }
         pipe_dtypes.update(new_bool_cols)
         pipe.dtypes = pipe_dtypes
@@ -2788,7 +2749,6 @@ def pipe_exists(
         debug=debug,
     )
     if debug:
-        from meerschaum.utils.debug import dprint
         dprint(f"{pipe} " + ('exists.' if exists else 'does not exist.'))
     return exists
 
@@ -3125,11 +3085,17 @@ def get_pipe_columns_types(
             debug=debug,
         )
 
+    if debug:
+        dprint(f"Fetching columns_types for {pipe} with via SQLAlchemy table.")
+
     table_columns = {}
     try:
         pipe_table = self.get_pipe_table(pipe, debug=debug)
         if pipe_table is None:
             return {}
+        if debug:
+            dprint(f"Found columns:")
+            mrsm.pprint(dict(pipe_table.columns))
         for col in pipe_table.columns:
             table_columns[str(col.name)] = str(col.type)
     except Exception as e:
@@ -3321,10 +3287,9 @@ def get_alter_columns_queries(
     -------
     A list of the `ALTER TABLE` SQL query or queries to be executed on the provided connector.
     """
-    if not pipe.exists(debug=debug):
+    if not pipe.exists(debug=debug) or pipe.static:
         return []
-    if pipe.static:
-        return
+
     from meerschaum.utils.sql import (
         sql_item_name,
         get_table_cols_types,
@@ -3370,7 +3335,8 @@ def get_alter_columns_queries(
             debug=debug,
         ).items()
     }
-    pipe_bool_cols = [col for col, typ in pipe.dtypes.items() if are_dtypes_equal(str(typ), 'bool')]
+    pipe_dtypes = pipe.dtypes
+    pipe_bool_cols = [col for col, typ in pipe_dtypes.items() if are_dtypes_equal(str(typ), 'bool')]
     pd_db_df_aliases = {
         'int': 'bool',
         'float': 'bool',
@@ -3378,7 +3344,10 @@ def get_alter_columns_queries(
         'guid': 'object',
     }
     if self.flavor == 'oracle':
-        pd_db_df_aliases['int'] = 'numeric'
+        pd_db_df_aliases.update({
+            'int': 'numeric',
+            'date': 'datetime',
+        })
 
     altered_cols = {
         col: (db_cols_types.get(col, 'object'), typ)
@@ -3387,6 +3356,10 @@ def get_alter_columns_queries(
         and not are_dtypes_equal(db_cols_types.get(col, 'object'), 'string')
     }
 
+    if debug and altered_cols:
+        dprint(f"Columns to be altered:")
+        mrsm.pprint(altered_cols)
+
     ### NOTE: Sometimes bools are coerced into ints or floats.
     altered_cols_to_ignore = set()
     for col, (db_typ, df_typ) in altered_cols.items():
@@ -3413,13 +3386,20 @@ def get_alter_columns_queries(
         if db_is_bool_compatible and df_is_bool_compatible:
             altered_cols_to_ignore.add(bool_col)
 
+    if debug and altered_cols_to_ignore:
+        dprint(f"Ignoring the following altered columns (false positives).")
+        mrsm.pprint(altered_cols_to_ignore)
+
     for col in altered_cols_to_ignore:
         _ = altered_cols.pop(col, None)
+
     if not altered_cols:
         return []
 
     if numeric_cols:
-        pipe.dtypes.update({col: 'numeric' for col in numeric_cols})
+        explicit_pipe_dtypes = pipe.get_dtypes(infer=False, debug=debug)
+        explicit_pipe_dtypes.update({col: 'numeric' for col in numeric_cols})
+        pipe.dtypes = explicit_pipe_dtypes
         if not pipe.temporary:
             edit_success, edit_msg = pipe.edit(debug=debug)
             if not edit_success:
@@ -3428,7 +3408,7 @@ def get_alter_columns_queries(
                     + f"{edit_msg}"
                 )
     else:
-        numeric_cols.extend([col for col, typ in pipe.dtypes.items() if typ.startswith('numeric')])
+        numeric_cols.extend([col for col, typ in pipe_dtypes.items() if typ.startswith('numeric')])
 
     numeric_type = get_db_type_from_pd_type('numeric', self.flavor, as_sqlalchemy=False)
     text_type = get_db_type_from_pd_type('str', self.flavor, as_sqlalchemy=False)
@@ -3636,20 +3616,18 @@ def get_to_sql_dtype(
     >>> get_to_sql_dtype(pipe, df)
     {'a': <class 'sqlalchemy.sql.sqltypes.JSON'>}
     """
-    from meerschaum.utils.dataframe import get_json_cols, get_numeric_cols, get_uuid_cols
+    from meerschaum.utils.dataframe import get_special_cols
     from meerschaum.utils.dtypes.sql import get_db_type_from_pd_type
     df_dtypes = {
         col: str(typ)
         for col, typ in df.dtypes.items()
     }
-    json_cols = get_json_cols(df)
-    numeric_cols = get_numeric_cols(df)
-    uuid_cols = get_uuid_cols(df)
-    df_dtypes.update({col: 'json' for col in json_cols})
-    df_dtypes.update({col: 'numeric' for col in numeric_cols})
-    df_dtypes.update({col: 'uuid' for col in uuid_cols})
+    special_cols = get_special_cols(df)
+    df_dtypes.update(special_cols)
+
     if update_dtypes:
         df_dtypes.update(pipe.dtypes)
+
     return {
         col: get_db_type_from_pd_type(typ, self.flavor, as_sqlalchemy=True)
         for col, typ in df_dtypes.items()
@@ -3920,3 +3898,15 @@ def get_temporary_target(
         + transact_id
         + ((separator + label) if label else '')
     )
+
+
+def _enforce_pipe_dtypes_chunks_hook(
+    pipe: mrsm.Pipe,
+    chunk_df: 'pd.DataFrame',
+    debug: bool = False,
+    **kwargs
+) -> 'pd.DataFrame':
+    """
+    Enforce a pipe's dtypes on each chunk.
+    """
+    return pipe.enforce_dtypes(chunk_df, debug=debug)