Skip to content

Commit

Permalink
changes for export
Browse files Browse the repository at this point in the history
  • Loading branch information
dzaslavskiy committed Jun 13, 2024
1 parent 8dac638 commit 18fdbc4
Show file tree
Hide file tree
Showing 8 changed files with 407 additions and 75 deletions.
6 changes: 3 additions & 3 deletions .github/workflows/deploy-app.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,6 @@ jobs:

- name: Add network policies
run: |
cf add-network-policy gdrive outbound-proxy -s ${{ steps.cf-setup.outputs.target-environment }}-public --protocol tcp --port 8080
cf add-network-policy gdrive es-proxy -s ${{ steps.cf-setup.outputs.target-environment }} --protocol tcp --port 8080
cf add-network-policy gdrive qualtrix -s ${{ steps.cf-setup.outputs.target-environment }} --protocol tcp --port 8080
cf add-network-policy util outbound-proxy -s prod-public --protocol tcp --port 8080
cf add-network-policy util es-proxy -s prod --protocol tcp --port 8080
cf add-network-policy util qualtrix -s prod --protocol tcp --port 8080
43 changes: 36 additions & 7 deletions gdrive/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,13 +53,18 @@ async def upload_file(
parent = drive_client.create_folder(id, settings.ROOT_DIRECTORY)

if zip:
with zipfile.ZipFile(stream) as archive:
files = archive.filelist
for file in files:
image = io.BytesIO(archive.read(file))
drive_client.upload_basic(
f"{filename}_{file.filename}", parent, image
)
try:
with zipfile.ZipFile(stream) as archive:
files = archive.filelist
for file in files:
image = io.BytesIO(archive.read(file))
client.upload_basic(
f"{filename}_{file.filename}", parent, image
)
except zipfile.BadZipFile as error:
client.upload_basic(filename, parent, stream)
log.error(f"An error occurred: {error}")
response.status_code = status.HTTP_400_BAD_REQUEST
else:
drive_client.upload_basic(filename, parent, stream)

Expand All @@ -85,3 +90,27 @@ async def delete_file(filename, response: Response):
except HttpError as error:
log.error(f"An error occurred: {error}")
response.status_code = error.status_code


@router.get("/list")
async def list():
""" """
s = client.list(count=200, parent="0AFrX3czp_UwZUk9PVA")
print(s)


@router.get("/count")
async def list():
""" """
s = drive_client.count(parent="1dz8UklyVsBDLP0wC5HPVSOuKpYGqrNEI")
print(s)


@router.post("/move")
async def move_file(file: str, source: str, dest: str):
return drive_client.move(file, source, dest)


@router.post("/create_folder")
async def move_file(name: str, dest: str):
return drive_client.create_folder(name, dest)
180 changes: 159 additions & 21 deletions gdrive/drive_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@
import logging
import json
import mimetypes
import time
from typing import List

from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseUpload
from googleapiclient.errors import HttpError

from gdrive import settings, error

Expand All @@ -30,14 +32,17 @@ def init():
log.info(f"Connected to Root Directory {driveId}")


def list(count: int = 10, shared: bool = True) -> None:
def list(count: int = 10, shared: bool = True, parent: str | None = None) -> None:
"""
Prints the names and ids of the first <count> files the user has access to.
"""

mq = f"'{parent}' in parents" if parent else ""

results = (
service.files()
.list(
q=mq,
pageSize=count,
fields="*",
supportsAllDrives=shared,
Expand All @@ -63,6 +68,37 @@ def list(count: int = 10, shared: bool = True) -> None:
log.info(f"No such key: {error} in {item}")


def count(shared: bool = True, parent: str | None = None) -> None:
"""
Prints the names and ids of the first <count> files the user has access to.
"""

mq = f"'{parent}' in parents" if parent else ""

results = __count(mq, "", shared)
count = len(results["files"])

while results.get("nextPageToken", None):
results = __count(mq, results["nextPageToken"], shared)
count += len(results["files"])
return count


def __count(query, nextPageToken, shared):
results = (
service.files()
.list(
q=query,
pageSize=1000,
supportsAllDrives=shared,
includeItemsFromAllDrives=shared,
pageToken=nextPageToken,
)
.execute()
)
return results


def create_empty_spreadsheet(filename: str, parent_id: str) -> str:
file_metadata = {
"name": filename,
Expand Down Expand Up @@ -180,17 +216,31 @@ def get_files_by_drive_id(filename: str, drive_id: str):
Get list of files by filename
"""

results = (
service.files()
.list(
q=f"name = '{filename}'",
corpora="drive",
driveId=drive_id,
includeTeamDriveItems=True,
supportsTeamDrives=True,
try:
results = (
service.files()
.list(
q=f"name = '{filename}'",
corpora="drive",
driveId=drive_id,
includeTeamDriveItems=True,
supportsTeamDrives=True,
)
.execute()
)
except HttpError: # manual retry hack
time.sleep(2)
results = (
service.files()
.list(
q=f"name = '{filename}'",
corpora="drive",
driveId=drive_id,
includeTeamDriveItems=True,
supportsTeamDrives=True,
)
.execute()
)
.execute()
)

return results["files"]

Expand All @@ -202,17 +252,35 @@ def get_files_in_folder(id: str) -> List:
files = []
page_token = None
while True:
results = (
service.files()
.list(
q=f"'{id}' in parents and trashed=false",
supportsAllDrives=True,
includeItemsFromAllDrives=True,
fields="nextPageToken, files(*)",
pageToken=page_token,
try:
results = (
service.files()
.list(
q=f"'{id}' in parents and trashed=false",
pageSize=1000,
supportsAllDrives=True,
includeItemsFromAllDrives=True,
fields="nextPageToken, files(*)",
pageToken=page_token,
orderBy="createdTime",
)
.execute()
)
except HttpError: # manual retry hack
time.sleep(2)
results = (
service.files()
.list(
q=f"'{id}' in parents and trashed=false",
pageSize=1000,
supportsAllDrives=True,
includeItemsFromAllDrives=True,
fields="nextPageToken, files(*)",
pageToken=page_token,
orderBy="createdTime",
)
.execute()
)
.execute()
)
files.extend(results.get("files", []))
page_token = results.get("nextPageToken")
if not page_token:
Expand All @@ -230,3 +298,73 @@ def delete_file(id: str) -> None:

def export(id: str) -> any:
return service.files().get_media(fileId=id).execute()


def copy(fileId: str, new_location: str):
new_file = {"parents": [new_location]}

result = (
service.files()
.copy(body=new_file, fileId=fileId, supportsAllDrives=True)
.execute()
)

return result


def search(responseId: str, source_drive: str):
"""
Prints the names and ids of the first <count> files the user has access to.
"""

mq = f"fullText contains '{responseId}' and name = 'analytics.json'"

results = (
service.files()
.list(
q=mq,
driveId=source_drive,
corpora="drive",
fields="*",
supportsAllDrives=True,
includeItemsFromAllDrives=True,
)
.execute()
)
items = results.get("files", [])
return items


def exists(filename: str, parent: str):
# , source_drive: str
"""
does file exist?
"""

results = (
service.files()
.list(
q=f"'{parent}' in parents and name = '{filename}'",
# driveId=source_drive,
# corpora='drive',
fields="*",
supportsAllDrives=True,
includeItemsFromAllDrives=True,
)
.execute()
)
# return results["files"]
items = results.get("files", [])
return items


def move(fileId: str, source: str, dest: str):
result = (
service.files()
.update(
addParents=dest, removeParents=source, fileId=fileId, supportsAllDrives=True
)
.execute()
)

return result
Loading

0 comments on commit 18fdbc4

Please sign in to comment.