Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add azure #68

Merged
merged 4 commits into from
Sep 20, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file.
66 changes: 66 additions & 0 deletions querent/collectors/azure/azure_collector.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
import asyncio
from typing import AsyncGenerator
import io

from azure.storage.blob import BlobServiceClient
from querent.config.collector_config import CollectorBackend, AzureCollectConfig
from querent.collectors.collector_base import Collector
from querent.collectors.collector_factory import CollectorFactory
from querent.collectors.collector_result import CollectorResult
from querent.common.uri import Uri


class AzureCollector(Collector):
def __init__(self, config: AzureCollectConfig, container_name: str, prefix: str):
self.account_url = config.account_url
self.blob_service_client = BlobServiceClient(
account_url=self.account_url, credential=config.credential
)
self.container_name = container_name
self.chunk_size = 1024
self.prefix = prefix

async def connect(self):
pass # No asynchronous connection needed for the Azure Blob Storage client

async def disconnect(self):
pass # No asynchronous disconnect needed for the Azure Blob Storage client

async def poll(self) -> AsyncGenerator[CollectorResult, None]:
container_client = self.blob_service_client.get_container_client(
self.container_name
)

async for blob in container_client.list_blobs(name_starts_with=self.prefix):
file = self.download_blob_as_byte_stream(container_client, blob.name)
async for chunk in self.read_chunks(file):
yield CollectorResult({"object_key": blob.name, "chunk": chunk})

async def read_chunks(self, file):
while True:
chunk = file.read(self.chunk_size)
if not chunk:
break
yield chunk

def download_blob_as_byte_stream(self, container_client, blob_name):
blob_client = container_client.get_blob_client(blob_name)
blob_properties = blob_client.get_blob_properties()
byte_stream = io.BytesIO()

if blob_properties["size"] > 0:
stream = blob_client.download_blob()
byte_stream.write(stream.readall())
byte_stream.seek(0) # Rewind the stream to the beginning

return byte_stream


class AzureCollectorFactory(CollectorFactory):
def backend(self) -> CollectorBackend:
return CollectorBackend.AzureBlobStorage

def resolve(self, uri: Uri, config: AzureCollectConfig) -> Collector:
container_name = uri.path.strip("/")
prefix = uri.query.get("prefix", "")
return AzureCollector(config, container_name, prefix)
4 changes: 4 additions & 0 deletions querent/config/collector_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ class FSCollectorConfig(BaseModel):
root_path: str
chunk_size: int = 1024

class AzureCollectConfig(BaseModel):
account_url: str
credential: str
chunk_size: int = 1024

class S3CollectConfig(BaseModel):
bucket: str
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -170,3 +170,4 @@ tika
openpyxl
coverage
pytest-cov
azure-storage-blob