snowflake-connector-python 3.16.0__cp313-cp313-win_amd64.whl

snowflake/connector/file_compression_type.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+from typing import NamedTuple
+
+
+class CompressionType(NamedTuple):
+    name: str
+    file_extension: str
+    mime_type: str
+    mime_subtypes: list[str]
+    is_supported: bool
+
+
+CompressionTypes = {
+    "GZIP": CompressionType(
+        name="GZIP",
+        file_extension=".gz",
+        mime_type="application",
+        mime_subtypes=["gzip", "x-gzip"],
+        is_supported=True,
+    ),
+    "DEFLATE": CompressionType(
+        name="DEFLATE",
+        file_extension=".deflate",
+        mime_type="application",
+        mime_subtypes=["zlib", "deflate"],
+        is_supported=True,
+    ),
+    "RAW_DEFLATE": CompressionType(
+        name="RAW_DEFLATE",
+        file_extension=".raw_deflate",
+        mime_type="application",
+        mime_subtypes=["raw_deflate"],
+        is_supported=True,
+    ),
+    "BZIP2": CompressionType(
+        name="BZIP2",
+        file_extension=".bz2",
+        mime_type="application",
+        mime_subtypes=["bzip2", "x-bzip2", "x-bz2", "x-bzip", "bz2"],
+        is_supported=True,
+    ),
+    "LZIP": CompressionType(
+        name="LZIP",
+        file_extension=".lz",
+        mime_type="application",
+        mime_subtypes=["lzip", "x-lzip"],
+        is_supported=False,
+    ),
+    "LZMA": CompressionType(
+        name="LZMA",
+        file_extension=".lzma",
+        mime_type="application",
+        mime_subtypes=["lzma", "x-lzma"],
+        is_supported=False,
+    ),
+    "LZO": CompressionType(
+        name="LZO",
+        file_extension=".lzo",
+        mime_type="application",
+        mime_subtypes=["lzo", "x-lzo"],
+        is_supported=False,
+    ),
+    "XZ": CompressionType(
+        name="XZ",
+        file_extension=".xz",
+        mime_type="application",
+        mime_subtypes=["xz", "x-xz"],
+        is_supported=False,
+    ),
+    "COMPRESS": CompressionType(
+        name="COMPRESS",
+        file_extension=".Z",
+        mime_type="application",
+        mime_subtypes=["compress", "x-compress"],
+        is_supported=False,
+    ),
+    "PARQUET": CompressionType(
+        name="PARQUET",
+        file_extension=".parquet",
+        mime_type="snowflake",
+        mime_subtypes=["parquet"],
+        is_supported=True,
+    ),
+    "ZSTD": CompressionType(
+        name="ZSTD",
+        file_extension=".zst",
+        mime_type="application",
+        mime_subtypes=["zstd", "x-zstd"],
+        is_supported=True,
+    ),
+    "BROTLI": CompressionType(
+        name="BROTLI",
+        file_extension=".br",
+        mime_type="application",
+        mime_subtypes=["br", "x-br"],
+        is_supported=True,
+    ),
+    "ORC": CompressionType(
+        name="ORC",
+        file_extension=".orc",
+        mime_type="snowflake",
+        mime_subtypes=["orc"],
+        is_supported=True,
+    ),
+}
+
+subtype_to_meta: dict[str, CompressionType] = {
+    ms.lower(): meta for meta in CompressionTypes.values() for ms in meta.mime_subtypes
+}
+
+# TODO: Snappy avro doesn't need to be compressed again
+
+
+def lookup_by_mime_sub_type(mime_subtype: str) -> CompressionType | None:
+    """Look up a CompressionType for a specific mime subtype."""
+    return subtype_to_meta.get(mime_subtype.lower())

snowflake/connector/file_lock.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import logging
+import time
+from os import stat_result
+from pathlib import Path
+from time import sleep
+
+MAX_RETRIES = 5
+INITIAL_BACKOFF_SECONDS = 0.025
+STALE_LOCK_AGE_SECONDS = 1
+
+
+class FileLockError(Exception):
+    pass
+
+
+class FileLock:
+    def __init__(self, path: Path) -> None:
+        self.path: Path = path
+        self.locked = False
+        self.logger = logging.getLogger(__name__)
+
+    def __enter__(self):
+        statinfo: stat_result | None = None
+        try:
+            statinfo = self.path.stat()
+        except FileNotFoundError:
+            pass
+        except OSError as e:
+            raise FileLockError(f"Failed to stat lock file {self.path} due to {e=}")
+
+        if statinfo and statinfo.st_ctime < time.time() - STALE_LOCK_AGE_SECONDS:
+            self.logger.debug("Removing stale file lock")
+            try:
+                self.path.rmdir()
+            except FileNotFoundError:
+                pass
+            except OSError as e:
+                raise FileLockError(
+                    f"Failed to remove stale lock file {self.path} due to {e=}"
+                )
+
+        backoff_seconds = INITIAL_BACKOFF_SECONDS
+        for attempt in range(MAX_RETRIES):
+            self.logger.debug(
+                f"Trying to acquire file lock after {backoff_seconds} seconds in attempt number {attempt}.",
+            )
+            backoff_seconds = backoff_seconds * 2
+            try:
+                self.path.mkdir(mode=0o700)
+                self.locked = True
+                break
+            except FileExistsError:
+                sleep(backoff_seconds)
+                continue
+            except OSError as e:
+                raise FileLockError(
+                    f"Failed to acquire lock file {self.path} due to {e=}"
+                )
+
+        if not self.locked:
+            raise FileLockError(
+                f"Failed to acquire file lock, after {MAX_RETRIES} attempts."
+            )
+
+    def __exit__(self, exc_type, exc_val, exc_tbc):
+        try:
+            self.path.rmdir()
+        except FileNotFoundError:
+            pass
+        self.locked = False