Merge pull request #103 from dpguthrie/sl-client

Semantic Layer Client

Author: dpguthrie

dbtc/client/admin.py

@@ -67,10 +67,8 @@ def wrapper(self, *args, **kwargs):
 
 
 class _AdminClient(_Client):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.session = requests.Session()
-        self.session.headers = self.headers
+    def __init__(self, session, **kwargs):
+        super().__init__(session, **kwargs)
 
     _path = None
 

dbtc/client/base.py

@@ -6,6 +6,7 @@
 from typing import Callable, Optional
 
 # third party
+import requests
 import rudder_analytics
 
 # first party
@@ -19,25 +20,34 @@
 class _Client(abc.ABC):
     def __init__(
         self,
+        session: requests.Session,
         *,
         api_key: str = None,
         service_token: str = None,
         host: str = None,
         do_not_track: bool = False,
+        environment_id: int = None,
+        use_beta_endpoint: bool = True,
     ):
-        self.api_key: Optional[str] = api_key or os.getenv('DBT_CLOUD_API_KEY', None)
+        self.api_key: Optional[str] = api_key or os.getenv("DBT_CLOUD_API_KEY", None)
         self.service_token: Optional[str] = service_token or os.getenv(
-            'DBT_CLOUD_SERVICE_TOKEN', None
+            "DBT_CLOUD_SERVICE_TOKEN", None
         )
         self._host: Optional[str] = host or os.getenv(
-            'DBT_CLOUD_HOST', self.DEFAULT_DOMAIN
+            "DBT_CLOUD_HOST", self.DEFAULT_DOMAIN
         )
+        self.environment_id = environment_id or os.getenv(
+            "DBT_CLOUD_ENVIRONMENT_ID", None
+        )
+        self._use_beta = use_beta_endpoint
         self.do_not_track: bool = do_not_track
         self._anonymous_id: str = str(uuid.uuid4())
         self._called_from: Optional[str] = None
         self.console = err_console
+        self.session = session
+        self.session.headers = self.headers
 
-    DEFAULT_DOMAIN = 'cloud.getdbt.com'
+    DEFAULT_DOMAIN = "cloud.getdbt.com"
 
     @property
     @abc.abstractmethod
@@ -51,29 +61,29 @@ def _header_property(self):
 
     @property
     def _base_url(self):
-        return f'https://{self._host}{self._path}'
+        return f"https://{self._host}{self._path}"
 
     @property
     def headers(self):
         return {
-            'Authorization': f'Token {getattr(self, self._header_property)}',
-            'Content-Type': 'application/json',
+            "Authorization": f"Bearer {getattr(self, self._header_property)}",
+            "Content-Type": "application/json",
         }
 
     def full_url(self, path: str = None):
         if path is not None:
-            return f'{self._base_url}{path}'
+            return f"{self._base_url}{path}"
 
         return self._base_url
 
     def _send_track(self, event_name: str, func: Callable, *args, **kwargs):
-        func_args = [a for a in inspect.getfullargspec(func).args if a != 'self']
+        func_args = [a for a in inspect.getfullargspec(func).args if a != "self"]
         properties = {
-            'method': func.__name__,
-            'dbtc_version': __version__,
-            'called_from': self._called_from,
+            "method": func.__name__,
+            "dbtc_version": __version__,
+            "called_from": self._called_from,
             **dict(zip(func_args, args)),
             **kwargs,
         }
-        properties = {k: v for k, v in properties.items() if not k.endswith('_id')}
+        properties = {k: v for k, v in properties.items() if not k.endswith("_id")}
         rudder_analytics.track(self._anonymous_id, event_name, properties)

dbtc/client/main.py

@@ -1,9 +1,15 @@
+# third party
+import requests
+
 # first party
 from dbtc.client.admin import _AdminClient
 from dbtc.client.metadata import _MetadataClient
+from dbtc.client.semantic_layer import _SemanticLayerClient
 
 
 class dbtCloudClient:
     def __init__(self, **kwargs):
-        self.cloud = _AdminClient(**kwargs)
-        self.metadata = _MetadataClient(**kwargs)
+        session = requests.Session()
+        self.cloud = _AdminClient(session, **kwargs)
+        self.metadata = _MetadataClient(session, **kwargs)
+        self.sl = _SemanticLayerClient(session, **kwargs)

dbtc/client/metadata.py

@@ -1,19 +1,13 @@
 # stdlib
 from typing import Any, Dict, List, Union
 
-# third party
-import requests
-
 # first party
 from dbtc.client.base import _Client
 
 
 class _MetadataClient(_Client):
-    def __init__(self, **kwargs):
-        self._use_beta = kwargs.pop("use_beta_endpoint", True)
-        super().__init__(**kwargs)
-        self.session = requests.Session()
-        self.session.headers = self.headers
+    def __init__(self, session, **kwargs):
+        super().__init__(session, **kwargs)
 
     _header_property = "service_token"