Merge pull request #477 from DrDroidLab/integration

Integration

.gitignore

@@ -174,4 +174,5 @@ dev_scripts
 # Ignore all local media files
 media/images/*
 media/csv_files/*
-media/text_files/*
+media/text_files/*
+env

accounts/cache.py

@@ -1,3 +1,4 @@
+import logging
 import traceback
 
 from django.conf import settings
@@ -6,6 +7,8 @@
 
 from accounts.models import AccountApiToken, Account
 
+logger = logging.getLogger(__name__)
+
 
 class AccountApiTokenCache:
     def __init__(self, cache_key, enabled):
@@ -79,7 +82,7 @@ def get(self, email):
             key = f'reset_password_token:{email}'
             return caches[self._cache_key].get(key)
         except Exception as ex:
-            print(
+            logger.error(
                 f"Error while getting reset password token from cache:: {traceback.format_exception(type(ex), ex, ex.__traceback__)}")
 
     def create_or_update(self, email, token, time_to_expire_s=None):
@@ -89,7 +92,7 @@ def create_or_update(self, email, token, time_to_expire_s=None):
             key = f'reset_password_token:{email}'
             caches[self._cache_key].set(key, token, time_to_expire_s)
         except Exception as ex:
-            print(
+            logger.error(
                 f"Error while setting reset password token in cache:: {traceback.format_exception(type(ex), ex, ex.__traceback__)}")
 
     def delete(self, email):
@@ -98,7 +101,7 @@ def delete(self, email):
             value = caches[self._cache_key].delete(key)
             return value
         except Exception as ex:
-            print(
+            logger.error(
                 f"Error while deleting reset password token from cache:: {traceback.format_exception(type(ex), ex, ex.__traceback__)}")
 
 

accounts/signals.py

@@ -1,3 +1,4 @@
+import logging
 from collections import Counter
 
 from allauth.account.signals import email_confirmed
@@ -10,6 +11,8 @@
 
 PUBLIC_EMAIL_DOMAINS = ['gmail.com', 'yahoo.com', 'hotmail.com', 'outlook.com', 'aol.com', 'icloud.com', 'zoho.com']
 
+logger = logging.getLogger(__name__)
+
 
 def get_account_for_user(instance: User) -> Account:
     user_email = instance.email
@@ -28,12 +31,12 @@ def get_account_for_user(instance: User) -> Account:
                 obj for obj in existing_same_domain_users if getattr(obj, 'account_id') == most_common_account_id)
             account = most_common_account_user.account
             found_associated_account = True
-            print(f'Found Associated account for user: {instance}')
+            logger.info(f'Found Associated account for user: {instance}')
 
     if not found_associated_account:
         account = Account(owner=instance)
         account.save()
-        print(f'Created account for user: {instance}')
+        logger.info(f'Created account for user: {instance}')
 
     return account
 
@@ -59,10 +62,10 @@ def create_user_account(sender, instance, created, **kwargs):
 
     instance.account = account
     instance.save()
-    print(f'Associated account for user: {instance}')
-    print(f'Created account setup task for: {instance}')
+    logger.info(f'Associated account for user: {instance}')
+    logger.info(f'Created account setup task for: {instance}')
     log_dict = {"msg": "Signup_New_User", "user_email": instance.email, "user_id": instance.id}
-    print(log_dict)
+    logger.info(log_dict)
 
 
 @receiver(email_confirmed)
@@ -73,6 +76,6 @@ def generate_account_token_on_email_confirmed(request, email_address, **kwargs):
         account = Account.objects.get(owner=user)
         token = AccountApiToken(account=account, created_by=user)
         token.save()
-        print(f'Created account api token for account: {account}')
+        logger.info(f'Created account api token for account: {account}')
     except Exception as e:
         pass

accounts/tasks.py

@@ -1,3 +1,4 @@
+import logging
 import uuid
 
 from celery import shared_task
@@ -8,6 +9,8 @@
 from django.core.mail import send_mail
 from django.template.loader import render_to_string
 
+logger = logging.getLogger(__name__)
+
 
 @shared_task
 def send_reset_password_email(email: str):
@@ -26,7 +29,7 @@ def send_reset_password_email(email: str):
             html_message=msg_html
         )
     except Exception as e:
-        print(str(e))
+        logger.error(str(e))
         pass
 
 
@@ -45,5 +48,5 @@ def send_user_invite_email(sender_name: str, invited_user_email: str, signup_dom
             html_message=msg_html
         )
     except Exception as e:
-        print(str(e))
+        logger.error(str(e))
         pass

accounts/views.py

@@ -1,3 +1,4 @@
+import logging
 import uuid
 from typing import Union
 
@@ -34,6 +35,8 @@
 from utils.proto_utils import proto_to_dict
 from utils.uri_utils import build_absolute_uri
 
+logger = logging.getLogger(__name__)
+
 
 @web_api(GetAccountApiTokensRequest)
 def get_account_api_tokens(request_message: GetAccountApiTokensRequest) -> \
@@ -265,7 +268,7 @@ def login_okta(request_message: HttpRequest) -> JsonResponse:
 
     if user:
         if not user.is_active:
-            print("Social Sign-in: User is not active - {}".format(user.id))
+            logger.error("Social Sign-in: User is not active - {}".format(user.id))
             return JsonResponse({'success': False, 'message': 'User is not active'}, status=403)
         user.last_login = timezone.now()
         user.save()
@@ -279,7 +282,7 @@ def login_okta(request_message: HttpRequest) -> JsonResponse:
         is_new_user = True
         EmailAddress.objects.create(user=user, email=email, verified=True, primary=True)
         generate_account_token_on_email_confirmed(request_message, email)
-        print("Social Sign-in: New user created - {}".format(user.id))
+        logger.info("Social Sign-in: New user created - {}".format(user.id))
 
     from dj_rest_auth.utils import jwt_encode
     auth_token, refresh_token = jwt_encode(user)

connectors/assets/extractor/azure_metadata_extractor.py

@@ -1,6 +1,11 @@
+import logging
+
 from connectors.assets.extractor.metadata_extractor import SourceMetadataExtractor
 from executor.source_processors.azure_api_processor import AzureApiProcessor
 from protos.base_pb2 import Source, SourceModelType
+from utils.logging_utils import log_function_call
+
+logger = logging.getLogger(__name__)
 
 
 class AzureConnectorMetadataExtractor(SourceMetadataExtractor):
@@ -9,12 +14,13 @@ def __init__(self, subscription_id, tenant_id, client_id, client_secret, account
         self.__client = AzureApiProcessor(subscription_id, tenant_id, client_id, client_secret)
         super().__init__(account_id, connector_id, Source.AZURE)
 
+    @log_function_call
     def extract_workspaces(self, save_to_db=False):
         model_type = SourceModelType.AZURE_WORKSPACE
         try:
             workspaces = self.__client.fetch_workspaces()
         except Exception as e:
-            print(f'Error fetching databases: {e}')
+            logger.error(f'Error fetching databases: {e}')
             return
         if not workspaces:
             return

connectors/assets/extractor/clickhouse_metadata_extractor.py

@@ -1,6 +1,11 @@
+import logging
+
 from connectors.assets.extractor.metadata_extractor import SourceMetadataExtractor
 from executor.source_processors.clickhouse_db_processor import ClickhouseDBProcessor
 from protos.base_pb2 import Source, SourceModelType
+from utils.logging_utils import log_function_call
+
+logger = logging.getLogger(__name__)
 
 
 class ClickhouseSourceMetadataExtractor(SourceMetadataExtractor):
@@ -10,12 +15,13 @@ def __init__(self, interface, host, port, user, password, account_id=None, conne
 
         super().__init__(account_id, connector_id, Source.CLICKHOUSE)
 
+    @log_function_call
     def extract_database(self, save_to_db=False):
         model_type = SourceModelType.CLICKHOUSE_DATABASE
         try:
             databases = self.__ch_db_processor.fetch_databases()
         except Exception as e:
-            print(f'Error fetching databases: {e}')
+            logger.error(f'Error fetching databases: {e}')
             return
         if not databases:
             return

connectors/assets/extractor/cloudwatch_metadata_extractor.py

@@ -1,6 +1,11 @@
+import logging
+
 from connectors.assets.extractor.metadata_extractor import SourceMetadataExtractor
 from executor.source_processors.aws_boto_3_api_processor import AWSBoto3ApiProcessor
 from protos.base_pb2 import Source, SourceModelType
+from utils.logging_utils import log_function_call
+
+logger = logging.getLogger(__name__)
 
 
 class CloudwatchSourceMetadataExtractor(SourceMetadataExtractor):
@@ -13,6 +18,7 @@ def __init__(self, region, aws_access_key=None, aws_secret_key=None, account_id=
 
         super().__init__(account_id, connector_id, Source.CLOUDWATCH)
 
+    @log_function_call
     def extract_metric(self, save_to_db=False):
         model_type = SourceModelType.CLOUDWATCH_METRIC
         model_data = {}
@@ -37,9 +43,10 @@ def extract_metric(self, save_to_db=False):
                 if save_to_db:
                     self.create_or_update_model_metadata(model_type, namespace, model_data[namespace])
         except Exception as e:
-            print(f'Error extracting metrics: {e}')
+            logger.error(f'Error extracting metrics: {e}')
         return model_data
 
+    @log_function_call
     def extract_log_groups(self, save_to_db=False):
         model_type = SourceModelType.CLOUDWATCH_LOG_GROUP
         model_data = {}
@@ -51,5 +58,5 @@ def extract_log_groups(self, save_to_db=False):
             if save_to_db:
                 self.create_or_update_model_metadata(model_type, self.__region, model_data[self.__region])
         except Exception as e:
-            print(f'Error extracting log groups: {e}')
+            logger.error(f'Error extracting log groups: {e}')
         return model_data