Time dagster asset loading (#4669)

* update deps

* chore: update oso_core deps

* feat(oso_core): use structured logging from structlog

* fix oso_core logging

* fix: allow for legacy definition import timing

* fix(oso_dagster): allow early asset factories to provide useful metadata

* feat(oso_dagster): asset loading

* fix(oso_dagster): decorated functions should use functools.wraps

* chore(ops): use legacy definitions file

* fix env vars

* improve docs

* fix some old references

* fix(oso_dagster): add some logging when using deferred table loading

* add some docs

* small clean up

* doc fix

* doc fix

Author: ravenac95

lib/oso-core/oso_core/logging/__init__.py

@@ -1,13 +1,13 @@
 import logging
 import os
-import sys
 import typing as t
 
-import colorlog
+import structlog
+from oso_core.logging.defaults import configure_structured_logging
 
 connected_to_application_logs = False
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 
 def add_oso_core_to_current_application_logging():
@@ -69,28 +69,9 @@ def setup_multiple_modules_logging(module_names: t.List[str]):
 def setup_module_logging(
     module_name: str,
     level: int = logging.DEBUG,
-    override_format: str = "",
     color: bool = False,
 ):
+    configure_structured_logging()
+
     logger = logging.getLogger(module_name)
     logger.setLevel(level)  # Adjust the level as needed
-
-    # Create a handler that logs to stdout
-    if color:
-        format = "%(asctime)s - %(log_color)s%(levelname)-8s%(reset)s - %(name)s - %(message)s"
-        stdout_handler = colorlog.StreamHandler(sys.stdout)
-        formatter = colorlog.ColoredFormatter(format, datefmt="%Y-%m-%dT%H:%M:%S")
-    else:
-        format = "%(asctime)s - %(levelname)-8s - %(name)s - %(message)s"
-        stdout_handler = logging.StreamHandler(sys.stdout)
-        formatter = logging.Formatter(format, datefmt="%Y-%m-%dT%H:%M:%S")
-    stdout_handler.setLevel(level)  # Adjust the level as needed
-
-    # Add the filter to the handler
-    stdout_handler.addFilter(ModuleFilter(module_name))
-
-    # Set a formatter (optional)
-    stdout_handler.setFormatter(formatter)
-
-    # Add the handler to the logger
-    logger.addHandler(stdout_handler)

lib/oso-core/oso_core/logging/decorators.py

@@ -6,10 +6,12 @@
 import logging
 
 
-def time_function(logger: logging.Logger):
+def time_function(logger: logging.Logger, override_name: str = ""):
     """Decorator to time a function and log the duration."""
 
     def decorator(func):
+        name = override_name or func.__name__
+
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             import time
@@ -18,7 +20,11 @@ def wrapper(*args, **kwargs):
             result = func(*args, **kwargs)
             end_time = time.time()
             duration = end_time - start_time
-            logger.info(f"Function {func.__name__} took {duration:.4f} seconds")
+            duration_ms = duration * 1000
+            logger.info(
+                f"Function `{name}` took {duration_ms:.4f} milliseconds.",
+                extra=dict(event_type="timing", duration_ms=duration_ms),
+            )
             return result
 
         return wrapper

lib/oso-core/oso_core/logging/defaults.py

@@ -0,0 +1,84 @@
+import logging
+
+import structlog
+
+STRUCTURED_LOGGING_CONFIGURED = False
+
+
+def flatten_extra_field(
+    logger: logging.Logger, log_method: str, event_dict: structlog.types.EventDict
+):
+    """This flattens the `extra` dict that is allowed in stdlib's logging so
+    that it appears as top level items in the log message."""
+    if "extra" in event_dict:
+        for key, value in event_dict["extra"].items():
+            event_dict[key] = value
+        del event_dict["extra"]
+    return event_dict
+
+
+def ensure_event_type(
+    logger: logging.Logger, log_method: str, event_dict: structlog.types.EventDict
+):
+    """This allows us to add an event_type to logs so we can do filtering in the
+    log UI"""
+
+    if "event_type" not in event_dict:
+        event_dict["event_type"] = "default"
+
+    return event_dict
+
+
+def configure_structured_logging() -> None:
+    """Configure structured logging for the application. This was taken almost
+    directly from the structlog docs. See:
+
+    https://www.structlog.org/en/stable/standard-library.html#rendering-using-structlog-based-formatters-within-logging
+
+    This sets up a structured handler at the root logger. It does not configure
+    the level of the root logger, so it is up to the user to enable the
+    appropriate level so that logs appear on stdout.
+
+    This should be called at the start of an application.
+    """
+    global STRUCTURED_LOGGING_CONFIGURED
+    if STRUCTURED_LOGGING_CONFIGURED:
+        return
+    STRUCTURED_LOGGING_CONFIGURED = True
+
+    timestamper = structlog.processors.TimeStamper(fmt="iso", utc=True)
+    shared_processors = [
+        flatten_extra_field,
+        structlog.stdlib.add_logger_name,
+        structlog.stdlib.add_log_level,
+        timestamper,
+        structlog.processors.StackInfoRenderer(),
+        ensure_event_type,
+    ]
+
+    structlog.configure(
+        processors=shared_processors
+        + [
+            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
+        ],
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        cache_logger_on_first_use=True,
+    )
+
+    formatter = structlog.stdlib.ProcessorFormatter(
+        # These run ONLY on `logging` entries that do NOT originate within
+        # structlog.
+        foreign_pre_chain=shared_processors,
+        # These run on ALL entries after the pre_chain is done.
+        processors=[
+            # Remove _record & _from_structlog.
+            structlog.stdlib.ProcessorFormatter.remove_processors_meta,
+            structlog.processors.JSONRenderer(),
+        ],
+    )
+
+    handler = logging.StreamHandler()
+    # Use OUR `ProcessorFormatter` to format all `logging` entries.
+    handler.setFormatter(formatter)
+    root_logger = logging.getLogger()
+    root_logger.addHandler(handler)

lib/oso-core/oso_core/logging/utils.py

@@ -0,0 +1,25 @@
+import logging
+from contextlib import contextmanager
+
+
+@contextmanager
+def time_context(logger: logging.Logger, name: str, **log_kwargs):
+    """Context manager to measure the timing of a block of code and log the
+    duration."""
+    import time
+
+    start_time = time.time()
+    try:
+        yield
+    finally:
+        end_time = time.time()
+        duration = end_time - start_time
+        duration_ms = duration * 1000
+        logger.info(
+            f"Timing `{name}` took {duration_ms:.4f} milliseconds.",
+            extra=dict(
+                event_type="timing",
+                duration_ms=duration_ms,
+                **log_kwargs,
+            ),
+        )

lib/oso-core/pyproject.toml

@@ -68,6 +68,7 @@ dependencies = [
     "pyarrow-stubs>=17.16,<21.0",
     "gcloud-aio-storage<10.0.0,>=9.3.0",
     "kr8s==0.20.9",
+    "structlog>=25.4.0",
 ]
 
 [build-system]

ops/helm-charts/oso-dagster/Chart.yaml

@@ -3,7 +3,7 @@ name: oso-dagster
 description: Extension of the dagster template
 
 type: application
-version: 0.19.0
+version: 0.20.0
 appVersion: "1.0.0"
 dependencies:
 - name: dagster

ops/helm-charts/oso-dagster/templates/config-map.yaml

@@ -25,6 +25,7 @@ data:
   DAGSTER_MCS_K8S_SERVICE_NAME: "{{ .Values.sqlmesh.mcs.k8s.serviceName }}"
   DAGSTER_MCS_K8S_NAMESPACE: "{{ .Values.sqlmesh.mcs.k8s.namespace }}"
   DAGSTER_SQLMESH_GATEWAY: "{{ .Values.sqlmesh.gateway }}"
+  DAGSTER_EAGERLY_LOAD_SQL_TABLES: "{{ .Values.osoDagster.eagerlyLoadSqlTables }}"
   SQLMESH_POSTGRES_INSTANCE_CONNECTION_STRING: "{{ .Values.configMap.secretPrefix }}-{{ .Values.secretmanagerKeys.sqlmeshPostgresInstanceStringConnectionString }}"
   SQLMESH_POSTGRES_USER: "{{ .Values.configMap.secretPrefix }}-{{ .Values.secretmanagerKeys.sqlmeshPostgresUser }}"
   SQLMESH_POSTGRES_PASSWORD: "{{ .Values.configMap.secretPrefix }}-{{ .Values.secretmanagerKeys.sqlmeshPostgresPassword }}"

ops/k8s-apps/base/dagster/dagster.yaml

@@ -88,7 +88,9 @@ spec:
             port: 3030
             dagsterApiGrpcArgs:
               - "-m"
-              - "oso_dagster.definitions"
+              - "oso_dagster.definitions.legacy"
+              - "--log-format"
+              - "json"
             includeConfigInLaunchedRuns:
               enabled: false
 

pyproject.toml

@@ -66,6 +66,7 @@ dependencies = [
     "pyarrow-stubs>=17.16,<21.0",
     "gcloud-aio-storage<10.0.0,>=9.3.0",
     "kr8s==0.20.9",
+    "structlog>=25.4.0",
 ]
 name = "oso"
 version = "1.0.0"
@@ -125,7 +126,7 @@ packages = [
 ]
 
 [tool.dagster]
-module_name = "oso_dagster.definitions"
+module_name = "oso_dagster.definitions.legacy"
 
 [tool.pyright]
 pythonVersion = "3.12"