traceable-agent 1.3.0__tar.gz → 2.0.0__tar.gz

{traceable-agent-1.3.0/src/traceable_agent.egg-info → traceable-agent-2.0.0}/PKG-INFO

@@ -1,17 +1,15 @@
 Metadata-Version: 2.1
 Name: traceable-agent
-Version: 1.3.0
+Version: 2.0.0
 Summary: Traceable.ai Python Agent
 Home-page: https://traceable.ai
 Author: Traceable.ai
+License: UNKNOWN
+Platform: UNKNOWN
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.7
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: hypertrace-agent==0.15.0
-Requires-Dist: psutil
-Requires-Dist: distro==1.6.0
-Requires-Dist: cffi
 
 # Python Agent
 
@@ -95,3 +93,5 @@ git submodule update --init --recursive
         ```
 - Test applications
     1. Refer to `docker/test_apps/`.
+
+

traceable-agent-2.0.0/setup.py

@@ -0,0 +1,67 @@
+# pylint: skip-file
+import os
+from setuptools import setup, find_packages, Extension
+
+exec(open('src/traceableai/version.py').read())
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setup(
+    name="traceable-agent",
+    version=__version__,
+    author="Traceable.ai",
+    description="Traceable.ai Python Agent",
+    url="https://traceable.ai",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Operating System :: OS Independent"
+    ],
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    include_package_data=True,
+    python_requires=">=3.8",
+    ext_modules=[
+        # this option will let wheel builder know to package c extensions when ran from CI. No effect for source tar installs
+        Extension(
+            name='traceableai.filter._libtraceable',
+            sources=[],
+            optional=os.environ.get('CIBUILDWHEEL', '0') != '1'
+        )
+    ],
+    install_requires=[
+        "opentelemetry-api==1.26.0",
+        "opentelemetry-exporter-otlp==1.26.0",
+        "opentelemetry-instrumentation==0.47b0",
+        "opentelemetry-instrumentation-aiohttp-client==0.47b0",
+        "opentelemetry-instrumentation-boto==0.47b0",
+        "opentelemetry-instrumentation-botocore==0.47b0",
+        "opentelemetry-instrumentation-wsgi==0.47b0",
+        "opentelemetry-instrumentation-fastapi==0.47b0",
+        "opentelemetry-instrumentation-flask==0.47b0",
+        "opentelemetry-instrumentation-mysql==0.47b0",
+        "opentelemetry-instrumentation-psycopg2==0.47b0",
+        "opentelemetry-instrumentation-requests==0.47b0",
+        "opentelemetry-instrumentation-grpc==0.47b0",
+        "opentelemetry-instrumentation-django==0.47b0",
+        "opentelemetry-instrumentation-aws-lambda==0.47b0",
+        "opentelemetry-propagator-b3==1.26.0",
+        "opentelemetry-proto==1.26.0",
+        "opentelemetry-sdk==1.26.0",
+        "opentelemetry-util-http==0.47b0",
+        "google>=3.0.0",
+        "pyyaml",
+        "protobuf<5",
+        "psutil",
+        "distro==1.6.0",
+        "cffi",
+        "setuptools",
+        "jaraco.text",
+        "platformdirs"
+    ],
+    entry_points={
+        'console_scripts': [
+            'traceableai-instrument = traceableai.autoinstrumentation.wrapper:run',
+        ],
+    }
+)

{traceable-agent-1.3.0 → traceable-agent-2.0.0/src/traceable_agent.egg-info}/PKG-INFO

@@ -1,17 +1,15 @@
 Metadata-Version: 2.1
 Name: traceable-agent
-Version: 1.3.0
+Version: 2.0.0
 Summary: Traceable.ai Python Agent
 Home-page: https://traceable.ai
 Author: Traceable.ai
+License: UNKNOWN
+Platform: UNKNOWN
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.7
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: hypertrace-agent==0.15.0
-Requires-Dist: psutil
-Requires-Dist: distro==1.6.0
-Requires-Dist: cffi
 
 # Python Agent
 
@@ -95,3 +93,5 @@ git submodule update --init --recursive
         ```
 - Test applications
     1. Refer to `docker/test_apps/`.
+
+

traceable-agent-2.0.0/src/traceable_agent.egg-info/SOURCES.txt

@@ -0,0 +1,45 @@
+MANIFEST.in
+README.md
+pyproject.toml
+setup.cfg
+setup.py
+src/traceable_agent.egg-info/PKG-INFO
+src/traceable_agent.egg-info/SOURCES.txt
+src/traceable_agent.egg-info/dependency_links.txt
+src/traceable_agent.egg-info/entry_points.txt
+src/traceable_agent.egg-info/requires.txt
+src/traceable_agent.egg-info/top_level.txt
+src/traceableai/__init__.py
+src/traceableai/agent.py
+src/traceableai/agent_init.py
+src/traceableai/constants.py
+src/traceableai/custom_logger.py
+src/traceableai/env.py
+src/traceableai/version.py
+src/traceableai/autoinstrumentation/__init__.py
+src/traceableai/autoinstrumentation/sitecustomize.py
+src/traceableai/autoinstrumentation/wrapper.py
+src/traceableai/config/__init__.py
+src/traceableai/config/config.py
+src/traceableai/config/config_pb2.py
+src/traceableai/config/default.py
+src/traceableai/config/environment.py
+src/traceableai/filter/__init__.py
+src/traceableai/filter/filter.py
+src/traceableai/filter/registry.py
+src/traceableai/filter/traceable.py
+src/traceableai/instrumentation/__init__.py
+src/traceableai/instrumentation/instrumentation_definitions.py
+src/traceableai/instrumentation/aiohttp/__init__.py
+src/traceableai/instrumentation/aws_lambda/__init__.py
+src/traceableai/instrumentation/boto/__init__.py
+src/traceableai/instrumentation/botocore/__init__.py
+src/traceableai/instrumentation/django/__init__.py
+src/traceableai/instrumentation/django/django_auto_instrumentation_compat.py
+src/traceableai/instrumentation/fast_api/__init__.py
+src/traceableai/instrumentation/fast_api/fast_api_auto_instrumentation_compat.py
+src/traceableai/instrumentation/flask/__init__.py
+src/traceableai/instrumentation/grpc/__init__.py
+src/traceableai/instrumentation/mysql/__init__.py
+src/traceableai/instrumentation/postgresql/__init__.py
+src/traceableai/instrumentation/requests/__init__.py

{traceable-agent-1.3.0 → traceable-agent-2.0.0}/src/traceable_agent.egg-info/entry_points.txt

@@ -1,2 +1,3 @@
 [console_scripts]
 traceableai-instrument = traceableai.autoinstrumentation.wrapper:run
+

traceable-agent-2.0.0/src/traceable_agent.egg-info/requires.txt

@@ -0,0 +1,28 @@
+opentelemetry-api==1.26.0
+opentelemetry-exporter-otlp==1.26.0
+opentelemetry-instrumentation==0.47b0
+opentelemetry-instrumentation-aiohttp-client==0.47b0
+opentelemetry-instrumentation-boto==0.47b0
+opentelemetry-instrumentation-botocore==0.47b0
+opentelemetry-instrumentation-wsgi==0.47b0
+opentelemetry-instrumentation-fastapi==0.47b0
+opentelemetry-instrumentation-flask==0.47b0
+opentelemetry-instrumentation-mysql==0.47b0
+opentelemetry-instrumentation-psycopg2==0.47b0
+opentelemetry-instrumentation-requests==0.47b0
+opentelemetry-instrumentation-grpc==0.47b0
+opentelemetry-instrumentation-django==0.47b0
+opentelemetry-instrumentation-aws-lambda==0.47b0
+opentelemetry-propagator-b3==1.26.0
+opentelemetry-proto==1.26.0
+opentelemetry-sdk==1.26.0
+opentelemetry-util-http==0.47b0
+google>=3.0.0
+pyyaml
+protobuf<5
+psutil
+distro==1.6.0
+cffi
+setuptools
+jaraco.text
+platformdirs

traceable-agent-2.0.0/src/traceableai/agent.py

@@ -0,0 +1,186 @@
+import os
+import sys
+import threading
+import traceback
+
+import distro
+
+from traceableai import constants
+from traceableai.agent_init import AgentInit
+from traceableai.config.config import Config
+from traceableai.custom_logger import get_custom_logger
+from traceableai.env import get_env_value
+from traceableai.filter.registry import Registry
+from traceableai.instrumentation.instrumentation_definitions import SUPPORTED_LIBRARIES, get_instrumentation_wrapper, \
+    FLASK_KEY, DJANGO_KEY, FAST_API_KEY, LAMBDA
+from traceableai.version import __version__ # pylint:disable=C0413
+logger = get_custom_logger(__name__)
+
+
+class Agent():
+    _instance = None
+    _singleton_lock = threading.Lock()
+
+    def __new__(cls):
+        '''constructor'''
+        if cls._instance is None:
+            with cls._singleton_lock:
+                logger.debug('Creating Agent')
+                logger.debug('Python version: %s', sys.version)
+                logger.debug('Traceable Agent version: %s', __version__)
+                cls._instance = super(Agent, cls).__new__(cls)
+        else:
+            logger.debug('Using existing Agent.')
+        return cls._instance
+
+    def __init__(self):
+        self._initialized = False
+        if not self._initialized:  # pylint: disable=E0203:
+            logger.debug('Initializing Agent.')
+            if not self.is_enabled():
+                return
+            try:
+                self._config = Config()
+                self._init = AgentInit(self._config)
+                self._initialized = True
+            except Exception as err:  # pylint: disable=W0703
+                logger.error('Failed to initialize Agent: exception=%s, stacktrace=%s',
+                             err,
+                             traceback.format_exc())
+        self.is_lambda = False
+        logger.debug("Platform: %s", distro.id())
+        logger.debug("Platform version: %s", distro.version())
+        logger.debug('TraceableAI Agent version: %s', __version__)
+        logger.debug("successfully initialized traceableai agent")
+
+        if hasattr(os, 'register_at_fork'):
+            logger.info('Registering after_in_child handler.')
+            os.register_at_fork(after_in_child=self.post_fork)  # pylint:disable=E1101
+
+    def post_fork(self):
+        logger.info("In post fork hook")
+        logger.info("Calling add traceable filter during post fork")
+        self._init.post_fork()
+        self.add_traceable_filter()
+
+    def instrument(self, app=None, skip_libraries=None, auto_instrument=False):
+        '''used to register applicable instrumentation wrappers'''
+        logger.debug("Beginning instrumentation")
+        self._init.apply_config(self._config)
+
+        if skip_libraries is None:
+            skip_libraries = []
+        if not self.is_initialized():
+            logger.debug('agent is not initialized, not instrumenting')
+            return
+
+        for library_key in SUPPORTED_LIBRARIES:
+            if library_key in skip_libraries:
+                logger.debug('not attempting to instrument %s', library_key)
+                continue
+            logger.debug("attempting to instrument %s", library_key)
+            self._instrument(library_key, app, auto_instrument)
+        logger.debug("Complete instrumentation")
+
+    def _instrument(self, library_key, app=None, auto_instrument=False):
+        """only used to allow the deprecated register_x library methods to still work"""
+
+        wrapper_instance = get_instrumentation_wrapper(library_key)
+        if wrapper_instance is None:
+            logger.debug("no instrumentation wrapper instance available for %s", library_key)
+            return
+
+        # Flask is a special case compared to rest of instrumentation
+        # using hypertrace-instrument we can replace flask class def before app is initialized
+        # however during code based instr we wrap the existing app
+        # since replacing class def after app is initialized doesnt have an effect
+        # the user has to pass the app in to agent.instrument()
+        # we could resolve this edge case by instead having users directly add the middleware
+        # ex: app = Flask();
+        # app = HypertraceMiddleware(App) => this in turn does agent.instrument()
+        # + we have ref to app
+        if library_key == FLASK_KEY and app is not None:
+            wrapper_instance.with_app(app)
+
+        # since ht sitecustomize pushes the agent to the front of the load path django instrumentation will error
+        # when using autoinstrumentation since we need the django app settings loaded to add middleware
+        #
+        # in order to resolve this if we detect django we wrap the wsgi/asgi app getter
+        # and instrument as soon as the app is retrieved(since settings have to be configured
+        # before returning loaded app)
+        if library_key == DJANGO_KEY and auto_instrument is True:
+            from traceableai.instrumentation.django.django_auto_instrumentation_compat import \
+                add_django_auto_instr_wrappers  # pylint: disable=C0415
+            add_django_auto_instr_wrappers(self, wrapper_instance)
+            return
+
+        # For FastAPI we need a handle to the user app before we can instrument fast & inject middleware
+        # to make things easier for the user always add the below instrumentation wrappers
+        # when FastAPI calls `.setup` take the app and then add instrumentation
+        if library_key == FAST_API_KEY:
+            from traceableai.instrumentation.fast_api.fast_api_auto_instrumentation_compat import \
+                add_fast_api_auto_instr_wrappers # pylint: disable=C0415
+            add_fast_api_auto_instr_wrappers(self, wrapper_instance)
+            return
+
+        if library_key == LAMBDA and '_HANDLER' not in os.environ:
+            return
+        logger.debug("registering library %s with wrapper instance", library_key)
+        self.register_library(library_key, wrapper_instance)
+
+
+    def register_library(self, library_name, wrapper_instance):
+        """will configure settings on an instrumentation wrapper + apply"""
+        logger.debug('attempting to register library instrumentation: %s', library_name)
+        try:
+            self._init.init_library_instrumentation(library_name, wrapper_instance)
+        except Exception as err:  # pylint: disable=W0703
+            logger.debug(constants.EXCEPTION_MESSAGE, library_name, err, traceback.format_exc())
+
+    def register_processor(self, processor) -> None:  # pylint: disable=R1710
+        '''Add additional span exporters + processors'''
+        logger.debug('Entering Agent.register_processor().')
+        logger.debug("initialized %s", self.is_initialized())
+        if not self.is_initialized():
+            return None
+        return self._init.register_processor(processor)
+
+
+
+    def is_enabled(self) -> bool:
+        '''Is agent enabled?'''
+        enabled = get_env_value('ENABLED')
+        if enabled:
+            if enabled.lower() == 'false':
+                logger.debug("ENABLED is disabled.")
+                return False
+        return True
+
+    def is_initialized(self) -> bool:
+        '''Is agent initialized - if an agent fails to init we should let the app continue'''
+        if not self.is_enabled():
+            return False
+        if not self._initialized:
+            return False
+        return True
+
+
+    def add_traceable_filter(self):
+        logger.debug("in add_traceable_filter")
+        if self.is_lambda:
+            logger.info('Not loading blocking extension - currently unsupported in lambda')
+            return
+        # We need to do a local import so that the extension is not loaded in a parent process
+        from traceableai.filter.traceable import Traceable, _LIBTRACEABLE_AVAILABLE  # pylint:disable=C0413,C0412,C0415
+        if not _LIBTRACEABLE_AVAILABLE:
+            logger.info("libtraceable unavailable, skipping filter registration")
+            return
+        if Config().config.blocking_config.enabled.value is not True:
+            logger.info("Not adding libtraceable filter - blocking is not enabled")
+            return
+        try:
+            Registry().register(Traceable)
+            logger.debug("successfully initialized traceable filter")
+        except Exception as exc: # pylint:disable=W0703
+            logger.debug(''.join(traceback.format_exception(None, exc, exc.__traceback__)))
+            logger.info("failed to register traceable filter")

traceable-agent-2.0.0/src/traceableai/agent_init.py

@@ -0,0 +1,177 @@
+'''Initialize all the components using configuration from AgentConfig'''
+import os
+import traceback
+import logging
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as OTLPGrpcSpanExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as OTLPHttpSpanExporter
+# from opentelemetry.exporter.zipkin.proto.http import ZipkinExporter
+from opentelemetry.trace import ProxyTracerProvider
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from traceableai import constants
+from traceableai.config import config_pb2
+
+# Initialize logger
+logger = logging.getLogger(__name__)  # pylint: disable=C0103
+
+
+class AgentInit:  # pylint: disable=R0902,R0903
+    '''Initialize all the OTel components using configuration from AgentConfig'''
+
+    def __init__(self, agent_config):
+        '''constructor'''
+        logger.debug('Initializing AgentInit object.')
+        self._config = agent_config
+
+        # Only available in python > 3.7
+        # this does prevent user from having to add post fork hooks to their
+        # web server config
+        if hasattr(os, 'register_at_fork'):
+            logger.info('Registering after_in_child handler.')
+            os.register_at_fork(after_in_child=self.post_fork)  # pylint:disable=E1101
+
+    def post_fork(self):
+        """Used to reinitialize exporter & processors in separate worker processes"""
+        self.apply_config(self._config)  # pylint:disable=W0212
+
+    def apply_config(self, agent_config):
+        """Initialize various aspects of the agent based on the most recent config"""
+        if agent_config:
+            self._config = agent_config
+
+        self.init_trace_provider()
+        self.init_propagation()
+
+        if "TA_ENABLE_CONSOLE_SPAN_EXPORTER" in os.environ:
+            self.set_console_span_processor()
+        else:
+            self.init_exporter()
+
+
+    def init_trace_provider(self) -> None:
+        '''Initialize trace provider and set resource attributes.'''
+        if isinstance(trace.get_tracer_provider(), ProxyTracerProvider):
+            logger.debug("no configured trace provider detected, adding one")
+            resource_attributes = {
+                "service.name": self._config.config.service_name.value,
+                "service.instance.id": os.getpid(),
+                "telemetry.sdk.version": constants.TELEMETRY_SDK_VERSION,
+                "telemetry.sdk.name": constants.TELEMETRY_SDK_NAME,
+                "telemetry.sdk.language": constants.TELEMETRY_SDK_LANGUAGE
+            }
+            if self._config.config.resource_attributes:
+                logger.debug(
+                    'Custom attributes found. Adding to resource attributes dict.')
+                resource_attributes.update(
+                    self._config.config.resource_attributes)
+            tracer_provider = TracerProvider(
+                resource=Resource.create(resource_attributes)
+            )
+            trace.set_tracer_provider(tracer_provider)
+        else:
+            logger.debug("tracer provider already configured, "
+                         "skipping trace provider configuration")
+
+    def init_exporter(self) -> None:
+        """Initialize exporter"""
+        reporter_type = self._config.config.reporting.trace_reporter_type
+        exporter = self._init_exporter(reporter_type)
+        if exporter is None:
+            logger.warning("Unable to initialize exporter")
+            return
+
+        span_processor = BatchSpanProcessor(exporter)
+        trace.get_tracer_provider().add_span_processor(span_processor)
+
+    def init_propagation(self) -> None:
+        '''Initialize requested context propagation protocols.'''
+        propagator_list = []
+        for prop_format in self._config.config.propagation_formats:
+            if prop_format == config_pb2.PropagationFormat.TRACECONTEXT:
+                from opentelemetry.trace.propagation.tracecontext \
+                  import TraceContextTextMapPropagator # pylint: disable=C0415
+                propagator_list += [ TraceContextTextMapPropagator() ]
+                logger.debug('Adding TRACECONTEXT trace propagator to list.')
+            if prop_format == config_pb2.PropagationFormat.B3:
+                from opentelemetry.propagators.b3 import B3MultiFormat  # pylint: disable=C0415
+                propagator_list += [ B3MultiFormat() ]
+                logger.debug('Adding B3 trace propagator to list.')
+
+        if len(propagator_list) == 0:
+            logger.debug('No propagators have been initialized.')
+
+        logger.debug('propagator_list: %s', str(propagator_list))
+        from opentelemetry.propagate import set_global_textmap # pylint: disable=C0415
+        from opentelemetry.propagators.composite import CompositePropagator # pylint: disable=C0415
+        composite_propagators = CompositePropagator(propagator_list)
+        set_global_textmap(composite_propagators)
+
+    def _set_wrapper_fields(self, wrapper):
+        data_cap = self._config.config.data_capture
+        wrapper.set_process_request_headers(data_cap.http_headers.request)
+        wrapper.set_process_request_body(data_cap.http_body.request)
+        wrapper.set_process_response_headers(data_cap.http_headers.response)
+        wrapper.set_process_response_body(data_cap.http_body.response)
+        wrapper.set_body_max_size(data_cap.body_max_size_bytes)
+
+    def init_library_instrumentation(self, instrumentation_name, wrapper_instance):
+        """used to configure instrumentation wrapper settings + apply instrumentation"""
+        logger.debug("Attempting to initialize %s instrumentation", instrumentation_name)
+        try:
+            self._set_wrapper_fields(wrapper_instance)
+            wrapper_instance.instrument()
+        except Exception as err: # pylint: disable=W0703
+            logger.debug(constants.INST_WRAP_EXCEPTION_MSSG,
+                         instrumentation_name,
+                         err,
+                         traceback.format_exc())
+
+    def register_processor(self, processor) -> None:
+        '''Register additional span exporter + processor'''
+        logger.debug('Entering AgentInit.register_processor().')
+        trace.get_tracer_provider().add_span_processor(processor)
+
+    def set_console_span_processor(self) -> None:
+        '''Register the console span processor for debugging purposes.'''
+        logger.debug('Entering AgentInit.setConsoleSpanProcessor().')
+        console_span_exporter = ConsoleSpanExporter(
+            service_name=self._config.config.service_name)
+        simple_export_span_processor = SimpleSpanProcessor(
+            console_span_exporter)
+        trace.get_tracer_provider().add_span_processor(simple_export_span_processor)
+
+    def _init_exporter(self, trace_reporter_type):
+        exporter_type = ''
+        exporter = None
+        try:
+            if trace_reporter_type == config_pb2.TraceReporterType.ZIPKIN:
+                exporter_type = 'zipkin'
+                #exporter = ZipkinExporter(
+                #    endpoint=self._config.agent_config.reporting.endpoint
+                #)
+            elif trace_reporter_type == config_pb2.TraceReporterType.OTLP:
+                exporter_type = 'otlp'
+                exporter = OTLPGrpcSpanExporter(endpoint=self._config.config.reporting.endpoint.value,
+                                            insecure= not self._config.config.reporting.secure.value)
+            elif trace_reporter_type == config_pb2.TraceReporterType.OTLP_HTTP:
+                exporter_type = 'otlp_http'
+                exporter = OTLPHttpSpanExporter(endpoint=self._config.config.reporting.endpoint.value)
+
+            if exporter:
+                logger.info('Initialized %s exporter reporting to `%s`',
+                            exporter_type,
+                            self._config.config.reporting.endpoint.value)
+            else:
+                logger.error("Unknown exporter type `%s`", trace_reporter_type)
+
+            return exporter
+        except Exception as err:  # pylint: disable=W0703
+            logger.error('Failed to initialize %s exporter: exception=%s, stacktrace=%s',
+                         exporter_type,
+                         err,
+                         traceback.format_exc())
+            return None

{traceable-agent-1.3.0 → traceable-agent-2.0.0}/src/traceableai/autoinstrumentation/sitecustomize.py

@@ -1,18 +1,10 @@
 import os
 
 import psutil
-# Allow use of TA_ env vars instead of HT_
-from hypertrace import env_var_settings
-
-from traceableai.hypertrace_env_override import add_deprecation_to_ht_env
-
-env_var_settings.ENV_VAR_PREFIXES.insert(0, 'TA')
-
-add_deprecation_to_ht_env()
 from traceableai.agent import Agent # pylint:disable=C0413
 
 from traceableai.config.config import Config  # pylint:disable=C0413,C0412
-from hypertrace.agent.custom_logger import get_custom_logger  # pylint:disable=C0413,C0412,C0411
+from traceableai.custom_logger import get_custom_logger  # pylint:disable=C0413,C0412,C0411
 
 config = Config()
 logger = get_custom_logger(__name__)

{traceable-agent-1.3.0 → traceable-agent-2.0.0}/src/traceableai/config/config.py

@@ -1,26 +1,27 @@
 import copy
 import json
-import logging
 import os
-from os import environ
 
 import yaml
 
 from google.protobuf import json_format
 
-from traceableai.config import traceable_config_pb2
+from traceableai.config import config_pb2
 from traceableai.config.default import DEFAULT
 from traceableai.config.environment import overwrite_with_environment
-logger = logging.getLogger(__name__)
+from traceableai.custom_logger import get_custom_logger
+logger = get_custom_logger(__name__)
 
 REMOTE_CONFIG_KEY = 'remote_config'
 BLOCKING_CONFIG_KEY = 'blocking_config'
 
 class Config(): # pylint:disable=R0903
+    _instance = None
+
     def __new__(cls):
-        if not hasattr(cls, 'instance'):
-            cls.instance = super(Config, cls).__new__(cls)
-        return cls.instance
+        if not hasattr(cls, "_instance") or cls._instance is None:
+            cls._instance = super(Config, cls).__new__(cls)
+        return cls._instance
 
     def __init__(self):
         if not hasattr(self, '_initialized'):
@@ -42,7 +43,7 @@ def merge_config(base_config, overriding_config):
 # 2.) Override the defaults with config file
 # 3.) Override the config file with env vars
 def build_config():
-    traceable_config = traceable_config_pb2.AgentConfig()
+    traceable_config = config_pb2.AgentConfig()
     config_dict = copy.deepcopy(DEFAULT)
     file_dict = read_from_file()
     if file_dict is not None:
@@ -62,7 +63,7 @@ def read_from_file():
         logger.debug("TA_CONFIG_FILE path not set")
         return None
 
-    with open(config_path, 'r') as config_file:
+    with open(config_path, 'r', encoding="UTF-8") as config_file:
         try:
             file_dict =  yaml.safe_load(config_file)
             # setting the new remote config with deprecated one if the new one doesn't exist
@@ -77,7 +78,7 @@ def read_from_file():
             return None
 
 def _config_file_path():
-    config_path = environ.get('TA_CONFIG_FILE', None)
+    config_path = os.environ.get('TA_CONFIG_FILE', None)
     if config_path is None:
         return None
     return config_path