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

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:

meerschaum/connectors/sql/tables/__init__.py

@@ -7,19 +7,24 @@ Define SQLAlchemy tables
 """
 
 from __future__ import annotations
+
+import pickle
+import threading
+import meerschaum as mrsm
 from meerschaum.utils.typing import Optional, Dict, Union, InstanceConnector, List
-from meerschaum.utils.warnings import error, warn
+from meerschaum.utils.warnings import error, warn, dprint
 
 ### store a tables dict for each connector
 connector_tables = {}
+_tables_locks = {}
 
 _sequence_flavors = {'duckdb', 'oracle'}
 _skip_index_names_flavors = {'mssql',}
 
 def get_tables(
     mrsm_instance: Optional[Union[str, InstanceConnector]] = None,
-    create: bool = True,
-    debug: Optional[bool] = None
+    create: Optional[bool] = None,
+    debug: bool = False,
 ) -> Union[Dict[str, 'sqlalchemy.Table'], bool]:
     """
     Create tables on the database and return the `sqlalchemy` tables.
@@ -29,10 +34,10 @@ def get_tables(
     mrsm_instance: Optional[Union[str, InstanceConnector]], default None
         The connector on which the tables reside.
 
-    create: bool, default True:
+    create: Optional[bool], default None
         If `True`, create the tables if they don't exist.
 
-    debug: Optional[bool], default None:
+    debug: bool, default False
         Verbosity Toggle.
 
     Returns
@@ -42,7 +47,6 @@ def get_tables(
 
     """
     from meerschaum.utils.debug import dprint
-    from meerschaum.utils.formatting import pprint
     from meerschaum.connectors.parse import parse_instance_keys
     from meerschaum.utils.packages import attempt_import
     from meerschaum.utils.sql import json_flavors
@@ -54,7 +58,7 @@ def get_tables(
         lazy=False,
     )
     if not sqlalchemy:
-        error(f"Failed to import sqlalchemy. Is sqlalchemy installed?")
+        error("Failed to import sqlalchemy. Is sqlalchemy installed?")
 
     if mrsm_instance is None:
         conn = get_connector(debug=debug)
@@ -63,149 +67,196 @@ def get_tables(
     else: ### NOTE: mrsm_instance MUST BE a SQL Connector for this to work!
         conn = mrsm_instance
 
-    ### kind of a hack. Create the tables remotely
-    from meerschaum.connectors.api import APIConnector
-    if isinstance(conn, APIConnector):
-        if create:
-            return conn.create_metadata(debug=debug)
-        return {}
+    cache_expired = _check_create_cache(conn, debug=debug) if conn.type == 'sql' else False
+    create = create or cache_expired
 
     ### Skip if the connector is not a SQL connector.
     if getattr(conn, 'type', None) != 'sql':
         return {}
 
-    if conn not in connector_tables:
-        if debug:
-            dprint(f"Creating tables for connector '{conn}'.")
-
-        id_type = sqlalchemy.Integer
-        if conn.flavor in json_flavors:
-            params_type = sqlalchemy.types.JSON
-        else:
-            params_type = sqlalchemy.types.Text
-        id_names = ('user_id', 'plugin_id', 'pipe_id')
-        sequences = {
-            k: sqlalchemy.Sequence(k + '_seq')
-            for k in id_names 
-        }
-        id_col_args = { k: [k, id_type] for k in id_names }
-        id_col_kw = { k: {'primary_key': True} for k in id_names }
-        index_names = conn.flavor not in _skip_index_names_flavors
-
-        if conn.flavor in _sequence_flavors:
-            for k, args in id_col_args.items():
-                args.append(sequences[k])
-            for k, kw in id_col_kw.items():
-                kw.update({'server_default': sequences[k].next_value()})
-
-        _tables = {
-            'users': sqlalchemy.Table(
-                'mrsm_users',
-                conn.metadata,
-                sqlalchemy.Column(
-                    *id_col_args['user_id'],
-                    **id_col_kw['user_id'],
+    conn_key = str(conn)
+    if conn_key in connector_tables:
+        return connector_tables[conn_key]
+
+    fasteners = attempt_import('fasteners')
+    pickle_path = conn.get_metadata_cache_path(kind='pkl')
+    lock_path = pickle_path.with_suffix('.lock')
+    lock = fasteners.InterProcessLock(lock_path)
+
+    with lock:
+        if not cache_expired and pickle_path.exists():
+            try:
+                with open(pickle_path, 'rb') as f:
+                    metadata = pickle.load(f)
+                metadata.bind = conn.engine
+                tables = {tbl.name.replace('mrsm_', ''): tbl for tbl in metadata.tables.values()}
+                connector_tables[conn_key] = tables
+                return tables
+            except Exception as e:
+                warn(f"Failed to load metadata from cache, rebuilding: {e}")
+
+        if conn_key not in _tables_locks:
+            _tables_locks[conn_key] = threading.Lock()
+
+    with _tables_locks[conn_key]:
+        if conn_key not in connector_tables:
+            if debug:
+                dprint(f"Building in-memory instance tables for '{conn}'.")
+
+            id_type = sqlalchemy.Integer
+            if conn.flavor in json_flavors:
+                from sqlalchemy.dialects.postgresql import JSONB
+                params_type = JSONB
+            else:
+                params_type = sqlalchemy.types.Text
+            id_names = ('user_id', 'plugin_id', 'pipe_id')
+            sequences = {
+                k: sqlalchemy.Sequence(k + '_seq')
+                for k in id_names 
+            }
+            id_col_args = { k: [k, id_type] for k in id_names }
+            id_col_kw = { k: {'primary_key': True} for k in id_names }
+            index_names = conn.flavor not in _skip_index_names_flavors
+
+            if conn.flavor in _sequence_flavors:
+                for k, args in id_col_args.items():
+                    args.append(sequences[k])
+                for k, kw in id_col_kw.items():
+                    kw.update({'server_default': sequences[k].next_value()})
+
+            _tables = {
+                'users': sqlalchemy.Table(
+                    'mrsm_users',
+                    conn.metadata,
+                    sqlalchemy.Column(
+                        *id_col_args['user_id'],
+                        **id_col_kw['user_id'],
+                    ),
+                    sqlalchemy.Column(
+                        'username',
+                        sqlalchemy.String(256),
+                        index = index_names,
+                        nullable = False,
+                    ),
+                    sqlalchemy.Column('password_hash', sqlalchemy.String(1024)),
+                    sqlalchemy.Column('email', sqlalchemy.String(256)),
+                    sqlalchemy.Column('user_type', sqlalchemy.String(256)),
+                    sqlalchemy.Column('attributes', params_type),
+                    extend_existing = True,
                 ),
-                sqlalchemy.Column(
-                    'username',
-                    sqlalchemy.String(256),
-                    index = index_names,
-                    nullable = False,
+                'plugins': sqlalchemy.Table(
+                    *([
+                        'mrsm_plugins',
+                        conn.metadata,
+                        sqlalchemy.Column(
+                            *id_col_args['plugin_id'],
+                            **id_col_kw['plugin_id'],
+                        ),
+                        sqlalchemy.Column(
+                            'plugin_name', sqlalchemy.String(256), index=index_names, nullable=False,
+                        ),
+                        sqlalchemy.Column('user_id', sqlalchemy.Integer, nullable=False),
+                        sqlalchemy.Column('version', sqlalchemy.String(256)),
+                        sqlalchemy.Column('attributes', params_type),
+                    ] + ([
+                        sqlalchemy.ForeignKeyConstraint(['user_id'], ['mrsm_users.user_id']),
+                    ] if conn.flavor != 'duckdb' else [])),
+                    extend_existing = True,
                 ),
-                sqlalchemy.Column('password_hash', sqlalchemy.String(1024)),
-                sqlalchemy.Column('email', sqlalchemy.String(256)),
-                sqlalchemy.Column('user_type', sqlalchemy.String(256)),
-                sqlalchemy.Column('attributes', params_type),
-                extend_existing = True,
-            ),
-            'plugins': sqlalchemy.Table(
-                *([
-                    'mrsm_plugins',
+                'temp_tables': sqlalchemy.Table(
+                    'mrsm_temp_tables',
                     conn.metadata,
                     sqlalchemy.Column(
-                        *id_col_args['plugin_id'],
-                        **id_col_kw['plugin_id'],
+                        'date_created',
+                        sqlalchemy.DateTime,
+                        index = True,
+                        nullable = False,
                     ),
                     sqlalchemy.Column(
-                        'plugin_name', sqlalchemy.String(256), index=index_names, nullable=False,
+                        'table',
+                        sqlalchemy.String(256),
+                        index = index_names,
+                        nullable = False,
                     ),
-                    sqlalchemy.Column('user_id', sqlalchemy.Integer, nullable=False),
-                    sqlalchemy.Column('version', sqlalchemy.String(256)),
-                    sqlalchemy.Column('attributes', params_type),
-                ] + ([
-                    sqlalchemy.ForeignKeyConstraint(['user_id'], ['mrsm_users.user_id']),
-                ] if conn.flavor != 'duckdb' else [])),
-                extend_existing = True,
-            ),
-            'temp_tables': sqlalchemy.Table(
-                'mrsm_temp_tables',
+                    sqlalchemy.Column(
+                        'ready_to_drop',
+                        sqlalchemy.DateTime,
+                        index = False,
+                        nullable = True,
+                    ),
+                    extend_existing = True,
+                ),
+            }
+
+            pipes_parameters_col = sqlalchemy.Column("parameters", params_type)
+            pipes_table_args = [
+                "mrsm_pipes",
                 conn.metadata,
                 sqlalchemy.Column(
-                    'date_created',
-                    sqlalchemy.DateTime,
-                    index = True,
+                    *id_col_args['pipe_id'],
+                    **id_col_kw['pipe_id'],
+                ),
+                sqlalchemy.Column(
+                    "connector_keys",
+                    sqlalchemy.String(256),
+                    index = index_names,
                     nullable = False,
                 ),
                 sqlalchemy.Column(
-                    'table',
+                    "metric_key",
                     sqlalchemy.String(256),
                     index = index_names,
                     nullable = False,
                 ),
                 sqlalchemy.Column(
-                    'ready_to_drop',
-                    sqlalchemy.DateTime,
-                    index = False,
+                    "location_key",
+                    sqlalchemy.String(256),
+                    index = index_names,
                     nullable = True,
                 ),
+                pipes_parameters_col,
+            ]
+            if conn.flavor in json_flavors:
+                pipes_table_args.append(
+                    sqlalchemy.Index(
+                        'ix_mrsm_pipes_parameters_tags',
+                        pipes_parameters_col['tags'],
+                        postgresql_using='gin'
+                    )
+                )
+            _tables['pipes'] = sqlalchemy.Table(
+                *pipes_table_args,
                 extend_existing = True,
-            ),
-        }
-
-        _tables['pipes'] = sqlalchemy.Table(
-            "mrsm_pipes",
-            conn.metadata,
-            sqlalchemy.Column(
-                *id_col_args['pipe_id'],
-                **id_col_kw['pipe_id'],
-            ),
-            sqlalchemy.Column(
-                "connector_keys",
-                sqlalchemy.String(256),
-                index = index_names,
-                nullable = False,
-            ),
-            sqlalchemy.Column(
-                "metric_key",
-                sqlalchemy.String(256),
-                index = index_names,
-                nullable = False,
-            ),
-            sqlalchemy.Column(
-                "location_key",
-                sqlalchemy.String(256),
-                index = index_names,
-                nullable = True,
-            ),
-            sqlalchemy.Column("parameters", params_type),
-            extend_existing = True,
-        )
-
-        ### store the table dict for reuse (per connector)
-        connector_tables[conn] = _tables
-        if create:
-            create_schemas(
-                conn,
-                schemas = [conn.internal_schema],
-                debug = debug,
             )
-            create_tables(conn, tables=_tables)
 
-    return connector_tables[conn]
+            ### store the table dict for reuse (per connector)
+            connector_tables[conn_key] = _tables
+
+            if debug:
+                dprint(f"Built in-memory tables for '{conn}'.")
+
+            if create:
+                if debug:
+                    dprint(f"Creating tables for connector '{conn}'.")
+
+                create_schemas(
+                    conn,
+                    schemas = [conn.internal_schema],
+                    debug = debug,
+                )
+                create_tables(conn, tables=_tables)
+
+                _write_create_cache(mrsm.get_connector(str(mrsm_instance)), debug=debug)
+
+        with open(pickle_path, 'wb') as f:
+            pickle.dump(conn.metadata, f)
+
+    connector_tables[conn_key] = _tables
+    return connector_tables[conn_key]
 
 
 def create_tables(
-    conn: 'meerschaum.connectors.SQLConnector',
+    conn: mrsm.connectors.SQLConnector,
     tables: Optional[Dict[str, 'sqlalchemy.Table']] = None,
 ) -> bool:
     """
@@ -224,14 +275,13 @@ def create_tables(
 
 
 def create_schemas(
-    conn: 'meerschaum.connectors.SQLConnector',
+    conn: mrsm.connectors.SQLConnector,
     schemas: List[str],
     debug: bool = False,
 ) -> bool:
     """
     Create the internal Meerschaum schema on the database.
     """
-    from meerschaum._internal.static import STATIC_CONFIG
     from meerschaum.utils.packages import attempt_import
     from meerschaum.utils.sql import sql_item_name, NO_SCHEMA_FLAVORS, SKIP_IF_EXISTS_FLAVORS
     if conn.flavor in NO_SCHEMA_FLAVORS:
@@ -257,3 +307,68 @@ def create_schemas(
             except Exception as e:
                 warn(f"Failed to create internal schema '{schema}':\n{e}")
     return all(successes.values())
+
+
+def _check_create_cache(connector: mrsm.connectors.SQLConnector, debug: bool = False) -> bool:
+    """
+    Return `True` if the metadata cache is missing or expired.
+    """
+    import json
+    from datetime import datetime, timedelta
+    from meerschaum.utils.dtypes import get_current_timestamp
+
+    if connector.type != 'sql':
+        return False
+
+    path = connector.get_metadata_cache_path()
+    if not path.exists():
+        if debug:
+            dprint(f"Metadata cache doesn't exist for '{connector}'.")
+        return True
+
+    try:
+        with open(path, 'r', encoding='utf-8') as f:
+            metadata = json.load(f)
+    except Exception:
+        return True
+
+    created_str = metadata.get('created', None)
+    if not created_str:
+        return True
+
+    now = get_current_timestamp()
+    created = datetime.fromisoformat(created_str)
+
+    delta = now - created
+    threshold_minutes = (
+        mrsm.get_config('system', 'connectors', 'sql', 'instance', 'create_metadata_cache_minutes')
+    )
+    threshold_delta = timedelta(minutes=threshold_minutes)
+    if delta >= threshold_delta:
+        if debug:
+            dprint(f"Metadata cache expired for '{connector}'.")
+        return True
+
+    if debug:
+        dprint(f"Using cached metadata for '{connector}'.")
+
+    return False
+
+
+def _write_create_cache(connector: mrsm.connectors.SQLConnector, debug: bool = False):
+    """
+    Write the current timestamp to the cache file.
+    """
+    if connector.type != 'sql':
+        return
+
+    import json
+    from meerschaum.utils.dtypes import get_current_timestamp, json_serialize_value
+
+    if debug:
+        dprint(f"Writing metadata cache for '{connector}'.")
+
+    path = connector.get_metadata_cache_path()
+    now = get_current_timestamp()
+    with open(path, 'w+', encoding='utf-8') as f:
+        json.dump({'created': now}, f, default=json_serialize_value)