AIDB Utility - going from unstructured data (local storage and aws s3) to blob table (#68)

* initial commit

* Basic Integration Working

* basic implementation finished

* test cases modified

* aws implementation working

* implementation finished

* remove images

* aidb utilities

* real images

* PR comments

* change type to text type

* pr comments

* remove data

* add datastore tests

* skip unit test

* image verify

Author: akash17mittal

.github/workflows/tests.yaml

@@ -41,3 +41,6 @@ jobs:
 
     - name: Run caching logic tests
       run: python3 -m tests.tests_caching_logic
+
+    - name: Run data store tests
+      run: python3 -m tests.tests_data_store

aidb/engine/base_engine.py

@@ -2,9 +2,6 @@
 from typing import Dict, List, Optional, Set, Tuple
 
 import pandas as pd
-import sqlalchemy
-import sqlalchemy.ext.asyncio
-import sqlalchemy.ext.automap
 
 from aidb.config.config import Config
 from aidb.config.config_types import InferenceBinding
@@ -14,7 +11,7 @@
 from aidb.query.query import FilteringClause, Query
 from aidb.query.utils import predicate_to_str
 from aidb.utils.asyncio import asyncio_run
-from aidb.utils.logger import logger
+from aidb.utils.db import infer_dialect, create_sql_engine
 
 
 class BaseEngine():
@@ -27,8 +24,8 @@ def __init__(
     self._connection_uri = connection_uri
     self._debug = debug
 
-    self._dialect = self._infer_dialect(connection_uri)
-    self._sql_engine = self._create_sql_engine()
+    self._dialect = infer_dialect(connection_uri)
+    self._sql_engine = create_sql_engine(connection_uri, debug)
 
     if infer_config:
       self._config: Config = asyncio_run(self._infer_config())
@@ -41,46 +38,6 @@ def __del__(self):
   # ---------------------
   # Setup
   # ---------------------
-  def _infer_dialect(self, connection_uri: str):
-    # Conection URIs have the following format:
-    # dialect+driver://username:password@host:port/database
-    # See https://docs.sqlalchemy.org/en/20/core/engines.html
-    dialect = connection_uri.split(':')[0]
-    if '+' in dialect:
-      dialect = dialect.split('+')[0]
-
-    supported_dialects = [
-      'mysql',
-      'postgresql',
-      'sqlite',
-    ]
-
-    if dialect not in supported_dialects:
-      logger.warning(
-        f'Unsupported dialect: {dialect}. Defaulting to mysql')
-      dialect = 'mysql'
-
-    return dialect
-
-
-  def _create_sql_engine(self):
-    logger.info(f'Creating SQL engine for {self._dialect}')
-    if self._dialect == 'mysql':
-      kwargs = {
-        'echo': self._debug,
-        'max_overflow': -1,
-      }
-    else:
-      kwargs = {}
-
-    engine = sqlalchemy.ext.asyncio.create_async_engine(
-      self._connection_uri,
-      **kwargs,
-    )
-
-    return engine
-
-
   async def _infer_config(self) -> Config:
     '''
     Infer the database configuration from the sql engine.

aidb/utils/db.py

@@ -0,0 +1,46 @@
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+import sqlalchemy.ext.automap
+
+from aidb.utils.logger import logger
+
+
+def infer_dialect(connection_uri: str):
+  # Conection URIs have the following format:
+  # dialect+driver://username:password@host:port/database
+  # See https://docs.sqlalchemy.org/en/20/core/engines.html
+  dialect = connection_uri.split(':')[0]
+  if '+' in dialect:
+    dialect = dialect.split('+')[0]
+
+  supported_dialects = [
+    'mysql',
+    'postgresql',
+    'sqlite',
+  ]
+
+  if dialect not in supported_dialects:
+    logger.warning(
+      f'Unsupported dialect: {dialect}. Defaulting to mysql')
+    dialect = 'mysql'
+
+  return dialect
+
+
+def create_sql_engine(connection_uri, debug=False):
+  dialect = infer_dialect(connection_uri)
+  logger.info(f'Creating SQL engine for {dialect}')
+  if dialect == 'mysql':
+    kwargs = {
+      'echo': debug,
+      'max_overflow': -1,
+    }
+  else:
+    kwargs = {}
+
+  engine = sqlalchemy.ext.asyncio.create_async_engine(
+    connection_uri,
+    **kwargs,
+  )
+
+  return engine

aidb_utilities/__init__.py

@@ -0,0 +1,53 @@
+from typing import List
+
+import boto3
+
+from aidb_utilities.blob_store.blob_store import Blob, BlobStore, DocumentBlob, ImageBlob
+from aidb_utilities.blob_store.utils import get_document_type, is_document, is_image_file
+
+
+class AwsS3BlobStore(BlobStore):
+
+  def __init__(self, bucket_name, access_key_id, secret_access_key):
+    session = boto3.Session(aws_access_key_id=access_key_id, aws_secret_access_key=secret_access_key)
+    s3 = session.resource('s3')
+    self.bucket = s3.Bucket(bucket_name)
+
+  def get_blobs(self) -> List[Blob]:
+    raise Exception("get_blobs function not implemented")
+
+
+class AwsS3ImageBlobStore(AwsS3BlobStore):
+
+  def __init__(self, bucket_name, access_key_id, secret_access_key):
+    super().__init__(bucket_name, access_key_id, secret_access_key)
+
+  def get_blobs(self) -> List[ImageBlob]:
+    image_blobs = []
+    image_count = 0
+    for obj in self.bucket.objects.all():
+      file_path = f"s3://{obj.bucket_name}/{obj.key}"
+      if is_image_file(file_path):
+        file_creation_time = str(obj.last_modified)
+        image_count += 1
+        image_blobs.append(ImageBlob(blob_id=image_count, image_path=file_path, created_at=file_creation_time))
+    return image_blobs
+
+
+class AwsS3DocumentBlobStore(AwsS3BlobStore):
+
+  def __init__(self, bucket_name, access_key_id, secret_access_key):
+    super().__init__(bucket_name, access_key_id, secret_access_key)
+
+  def get_blobs(self) -> List[DocumentBlob]:
+    doc_blobs = []
+    doc_count = 0
+    for obj in self.bucket.objects.all():
+      file_path = f"s3://{obj.bucket_name}/{obj.key}"
+      if is_document(file_path):
+        file_creation_time = str(obj.last_modified)
+        doc_type = get_document_type(file_path)
+        doc_count += 1
+        doc_blobs.append(
+          DocumentBlob(blob_id=doc_count, doc_path=file_path, created_at=file_creation_time, doc_type=doc_type))
+    return doc_blobs

aidb_utilities/blob_store/__init__.py

@@ -0,0 +1,61 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Union
+
+
+class DocumentType(Enum):
+  PDF = 'pdf'
+  DOCX = 'docx'
+  DOC = 'doc'
+
+
+@dataclass
+class Blob:
+  blob_id: int
+
+  def to_dict(self):
+    return {
+      'blob_id': self.blob_id,
+    }
+
+
+@dataclass
+class ImageBlob(Blob):
+  image_path: str
+  created_at: Union[None, str]
+
+  def to_dict(self):
+    return {
+      'blob_id': self.blob_id,
+      'image_path': self.image_path,
+      'created_at': self.created_at
+    }
+
+
+@dataclass
+class DocumentBlob(Blob):
+  doc_path: str
+  created_at: str
+  doc_type: DocumentType
+
+  def to_dict(self):
+    return {
+      'blob_id': self.blob_id,
+      'doc_path': self.doc_path,
+      'created_at': self.created_at,
+      'doc_type': self.doc_type.value
+    }
+
+
+class BlobStore(ABC):
+
+  def __int__(self):
+    '''
+    configuration, data store access keys, etc.
+    '''
+    pass
+
+  @abstractmethod
+  def get_blobs(self) -> List[Blob]:
+    raise Exception("get_blobs function not implemented")

aidb_utilities/blob_store/aws_s3_storage.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import List
+
+from aidb_utilities.blob_store.blob_store import Blob
+
+from aidb_utilities.blob_store.blob_store import BlobStore, DocumentBlob, ImageBlob
+from aidb_utilities.blob_store.utils import get_document_type, get_local_file_creation_time, is_document, is_image_file
+
+
+class LocalBlobStore(BlobStore):
+
+  def __init__(self, local_dir):
+    self.files = Path(local_dir).rglob('*')
+
+  def get_blobs(self) -> List[Blob]:
+    raise Exception("get_blobs function not implemented")
+
+
+class LocalImageBlobStore(LocalBlobStore):
+
+  def __init__(self, local_dir):
+    super().__init__(local_dir)
+
+  def get_blobs(self) -> List[ImageBlob]:
+    image_blobs = []
+    image_count = 0
+    for file in self.files:
+      file_path = str(file)
+      if is_image_file(file_path):
+        file_creation_time = get_local_file_creation_time(file_path)
+        image_count += 1
+        image_blobs.append(ImageBlob(blob_id=image_count, image_path=file_path, created_at=file_creation_time))
+    return image_blobs
+
+
+class LocalDocumentBlobStore(LocalBlobStore):
+
+  def __init__(self, local_dir):
+    super().__init__(local_dir)
+
+  def get_blobs(self) -> List[DocumentBlob]:
+    doc_blobs = []
+    doc_count = 0
+    for file in self.files:
+      file_path = str(file)
+      if is_document(file_path):
+        file_creation_time = get_local_file_creation_time(file_path)
+        doc_type = get_document_type(file_path)
+        doc_count += 1
+        doc_blobs.append(
+          DocumentBlob(blob_id=doc_count, doc_path=file_path, created_at=file_creation_time, doc_type=doc_type))
+    return doc_blobs

aidb_utilities/blob_store/blob_store.py

@@ -0,0 +1,40 @@
+import os
+import time
+
+from PIL import Image
+from aidb_utilities.blob_store.blob_store import DocumentType
+
+
+def get_file_extension(file: str):
+  return file.split('.')[-1]
+
+
+def get_local_file_creation_time(file: str):
+  return time.ctime(os.path.getctime(file))
+
+
+def is_image_file(file: str):
+  img = Image.open(file)
+  try:
+    img.verify()
+    return True
+  except Exception:
+    return False
+
+
+def is_document(file: str):
+  document_extension_filter = ['doc', 'docx', 'pdf']
+  file_ext = get_file_extension(file)
+  return file_ext in document_extension_filter
+
+
+def get_document_type(file: str):
+  file_ext = get_file_extension(file)
+  if file_ext == 'pdf':
+    return DocumentType.PDF
+  elif file_ext == 'doc':
+    return DocumentType.DOC
+  elif file_ext == 'docx':
+    return DocumentType.DOCX
+  else:
+    raise Exception("Unsupported document type")