From 7bd38d6ab5e8a6d8237fbe84d5221d8fdaa61e10 Mon Sep 17 00:00:00 2001
From: JMkrish <mural.krish4u@gmail.com>
Date: Thu, 28 Mar 2024 00:49:26 +0530
Subject: [PATCH] ENG-4971: Extending client object creation to support either
 a full URL or a tenant

---
 smsdk/client.py    | 15 +--------------
 smsdk/client_v0.py | 30 +++++++++++++++++++++++-------
 2 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/smsdk/client.py b/smsdk/client.py
index 4ac5778..cb383d6 100644
--- a/smsdk/client.py
+++ b/smsdk/client.py
@@ -125,10 +125,6 @@ def dict_to_df(data, normalize=True):
 class Client(ClientV0):
     """Connection point to the Sight Machine platform to retrieve data"""
 
-    session = None
-    tenant = None
-    config = None
-
     def __init__(
         self, tenant: str, site_domain: str = "sightmachine.io", protocol: str = "https"
     ):
@@ -143,16 +139,7 @@ def __init__(
         :type site_domain: :class:`string`
         """
 
-        self.tenant = tenant
-
-        # Handle internal configuration
-        self.config = {}
-        self.config["protocol"] = protocol
-        self.config["site.domain"] = site_domain
-
-        # Setup Authenticator
-        self.auth = Authenticator(self)
-        self.session = self.auth.session
+        super().__init__(tenant, site_domain=site_domain, protocol=protocol)
 
     @version_check_decorator
     def select_db_schema(self, schema_name):
diff --git a/smsdk/client_v0.py b/smsdk/client_v0.py
index 5dacf43..ef90c55 100644
--- a/smsdk/client_v0.py
+++ b/smsdk/client_v0.py
@@ -105,7 +105,9 @@ class ClientV0(object):
     tenant = None
     config = None
 
-    def __init__(self, tenant, site_domain="sightmachine.io"):
+    def __init__(
+        self, tenant: str, site_domain: str = "sightmachine.io", protocol: str = "https"
+    ):
         """
         Initialize the client.
 
@@ -117,12 +119,26 @@ def __init__(self, tenant, site_domain="sightmachine.io"):
         :type site_domain: :class:`string`
         """
 
-        self.tenant = tenant
-
-        # Handle internal configuration
-        self.config = {}
-        self.config["protocol"] = "https"
-        self.config["site.domain"] = site_domain
+        if "://" in tenant:  # Full URL provided eg. "https://demo.sightmachine.io"
+            protocol, domain = tenant.split("://")[0], tenant.split("://")[1]
+            self.tenant = domain.split(".")[0]
+            self.config = {
+                "protocol": protocol,
+                "site.domain": ".".join(domain.split(".")[-2:]),
+            }
+        else:
+            tenant_parts = tenant.split(".")
+            if (
+                len(tenant_parts) > 1
+            ):  # Tenant with domain provided but missing protocol eg. "demo.sightmachine.io"
+                self.tenant = tenant_parts[0]
+                self.config = {
+                    "protocol": protocol,
+                    "site.domain": ".".join(tenant_parts[1:]),
+                }
+            else:  # Only tenant name provided eg. "demo"
+                self.tenant = tenant
+                self.config = {"protocol": protocol, "site.domain": site_domain}
 
         # Setup Authenticator
         self.auth = Authenticator(self)