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

meerschaum/connectors/sql/_create_engine.py

@@ -31,7 +31,6 @@ install_flavor_drivers = {
     'mssql': ['pyodbc'],
     'oracle': ['oracledb'],
 }
-require_patching_flavors = {'cockroachdb': [('sqlalchemy-cockroachdb', 'sqlalchemy_cockroachdb')]}
 
 flavor_dialects = {
     'cockroachdb': (
@@ -63,19 +62,6 @@ def create_engine(
         )
         if self.flavor == 'mssql':
             _init_mssql_sqlalchemy()
-    if self.flavor in require_patching_flavors:
-        from meerschaum.utils.packages import determine_version, _monkey_patch_get_distribution
-        import pathlib
-        for install_name, import_name in require_patching_flavors[self.flavor]:
-            pkg = attempt_import(
-                import_name,
-                debug=debug,
-                lazy=False,
-                warn=False
-            )
-            _monkey_patch_get_distribution(
-                install_name, determine_version(pathlib.Path(pkg.__file__), venv='mrsm')
-            )
 
     ### supplement missing values with defaults (e.g. port number)
     for a, value in flavor_configs[self.flavor]['defaults'].items():
@@ -189,6 +175,9 @@ def _init_mssql_sqlalchemy():
         lazy=False,
         warn=False,
     )
+    if pyodbc is None:
+        raise EnvironmentError("Cannot import pyodbc. Is the MSSQL driver installed?")
+
     pyodbc.pooling = False
 
     MSDialect_pyodbc = sqlalchemy_dialects_mssql_pyodbc.MSDialect_pyodbc

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
 
@@ -170,7 +169,6 @@ def fetch_pipes_keys(
     debug: bool, default False
         Verbosity toggle.
     """
-    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
@@ -338,7 +336,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 +389,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 +1004,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 +1042,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 +1061,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 +1129,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 +1158,25 @@ def get_pipe_data(
         **kw
     )
 
+    read_kwargs = {}
     if is_dask:
         index_col = pipe.columns.get('datetime', None)
-        kw['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))
+        read_kwargs['index_col'] = index_col
 
-    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)
+    if as_iterator:
+        return chunks
 
-    for col in bytes_columns:
-        if col not in df.columns:
-            continue
-        df[col] = df[col].apply(attempt_cast_to_bytes)
-
-    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 +1509,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 +1538,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 +1690,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 +2714,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 +3050,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 +3252,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 +3300,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 +3309,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 +3321,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 +3351,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 +3373,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 +3581,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 +3863,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)

meerschaum/connectors/sql/_sql.py

@@ -131,23 +131,28 @@ def read(
     """
     if chunks is not None and chunks <= 0:
         return []
+
     from meerschaum.utils.sql import sql_item_name, truncate_item_name
     from meerschaum.utils.dtypes import are_dtypes_equal, coerce_timezone
     from meerschaum.utils.dtypes.sql import TIMEZONE_NAIVE_FLAVORS
     from meerschaum.utils.packages import attempt_import, import_pandas
     from meerschaum.utils.pool import get_pool
     from meerschaum.utils.dataframe import chunksize_to_npartitions, get_numeric_cols
+    from meerschaum.utils.misc import filter_arguments
     import warnings
     import traceback
     from decimal import Decimal
+
     pd = import_pandas()
     dd = None
+
     is_dask = 'dask' in pd.__name__
     pandas = attempt_import('pandas')
     is_dask = dd is not None
     npartitions = chunksize_to_npartitions(chunksize)
     if is_dask:
         chunksize = None
+
     schema = schema or self.schema
     utc_dt_cols = [
         col
@@ -158,7 +163,7 @@ def read(
     if dtype and utc_dt_cols and self.flavor in TIMEZONE_NAIVE_FLAVORS:
         dtype = dtype.copy()
         for col in utc_dt_cols:
-            dtype[col] = 'datetime64[ns]'
+            dtype[col] = 'datetime64[us]'
 
     pool = get_pool(workers=workers)
     sqlalchemy = attempt_import("sqlalchemy", lazy=False)
@@ -222,26 +227,33 @@ def read(
         else format_sql_query_for_dask(str_query)
     )
 
+    def _get_chunk_args_kwargs(_chunk):
+        return filter_arguments(
+            chunk_hook,
+            _chunk,
+            workers=workers,
+            chunksize=chunksize,
+            debug=debug,
+            **kw
+        )
+
     chunk_list = []
     chunk_hook_results = []
     def _process_chunk(_chunk, _retry_on_failure: bool = True):
         if self.flavor in TIMEZONE_NAIVE_FLAVORS:
             for col in utc_dt_cols:
-                _chunk[col] = coerce_timezone(_chunk[col], strip_timezone=False)
+                _chunk[col] = coerce_timezone(_chunk[col], strip_utc=False)
         if not as_hook_results:
             chunk_list.append(_chunk)
+
         if chunk_hook is None:
             return None
 
+        chunk_args, chunk_kwargs = _get_chunk_args_kwargs(_chunk)
+
         result = None
         try:
-            result = chunk_hook(
-                _chunk,
-                workers=workers,
-                chunksize=chunksize,
-                debug=debug,
-                **kw
-            )
+            result = chunk_hook(*chunk_args, **chunk_kwargs)
         except Exception:
             result = False, traceback.format_exc()
             from meerschaum.utils.formatting import get_console
@@ -292,8 +304,16 @@ def read(
                         self.engine,
                         **read_sql_query_kwargs
                     )
+
                     to_return = (
-                        chunk_generator
+                        (
+                            chunk_generator
+                            if not (as_hook_results or chunksize is None)
+                            else (
+                                _process_chunk(_chunk)
+                                for _chunk in chunk_generator
+                            )
+                        )
                         if as_iterator or chunksize is None
                         else (
                             list(pool.imap(_process_chunk, chunk_generator))
@@ -339,9 +359,8 @@ def read(
         try:
             for chunk in chunk_generator:
                 if chunk_hook is not None:
-                    chunk_hook_results.append(
-                        chunk_hook(chunk, chunksize=chunksize, debug=debug, **kw)
-                    )
+                    chunk_args, chunk_kwargs = _get_chunk_args_kwargs(chunk)
+                    chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
                 chunk_list.append(chunk)
                 read_chunks += 1
                 if chunks is not None and read_chunks >= chunks:
@@ -356,9 +375,8 @@ def read(
     try:
         for chunk in chunk_generator:
             if chunk_hook is not None:
-                chunk_hook_results.append(
-                    chunk_hook(chunk, chunksize=chunksize, debug=debug, **kw)
-                )
+                chunk_args, chunk_kwargs = _get_chunk_args_kwargs(chunk)
+                chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
             chunk_list.append(chunk)
             read_chunks += 1
             if chunks is not None and read_chunks >= chunks:
@@ -389,9 +407,8 @@ def read(
     ### call the hook on any missed chunks.
     if chunk_hook is not None and len(chunk_list) > len(chunk_hook_results):
         for c in chunk_list[len(chunk_hook_results):]:
-            chunk_hook_results.append(
-                chunk_hook(c, chunksize=chunksize, debug=debug, **kw)
-            )
+            chunk_args, chunk_kwargs = _get_chunk_args_kwargs(c)
+            chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
 
     ### chunksize is not None so must iterate
     if debug:
@@ -784,6 +801,7 @@ def to_sql(
     from meerschaum.utils.warnings import error, warn
     import warnings
     import functools
+    import traceback
 
     if name is None:
         error(f"Name must not be `None` to insert data into {self}.")
@@ -1057,7 +1075,7 @@ def to_sql(
     except Exception as e:
         if not silent:
             warn(str(e))
-        success, msg = False, str(e)
+        success, msg = False, traceback.format_exc()
 
     end = time.perf_counter()
     if success: