File size: 669 Bytes
e4f9cbe
55dc3dd
 
e4f9cbe
 
c14732f
e4f9cbe
 
55dc3dd
 
d9bd6c3
55dc3dd
d9bd6c3
 
e4f9cbe
 
 
c14732f
e4f9cbe
c14732f
 
 
e4f9cbe
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
"""Utils for duckdb."""
import os

import duckdb

from ..config import data_path, env


def duckdb_setup(con: duckdb.DuckDBPyConnection) -> str:
  """Setup DuckDB. This includes setting up the extensions directory and GCS access."""
  con.execute(f"""
    SET extension_directory='{os.path.join(data_path(), '.duckdb')}';
  """)

  con.install_extension('httpfs')
  con.load_extension('httpfs')

  if env('GCS_REGION'):
    return f"""
        SET s3_region='{env('GCS_REGION')}';
        SET s3_access_key_id='{env('GCS_ACCESS_KEY')}';
        SET s3_secret_access_key='{env('GCS_SECRET_KEY')}';
        SET s3_endpoint='storage.googleapis.com';
      """
  return ''