github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/actions/lua/storage/azure/blob.go (about)

     1  package azure
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"regexp"
     8  	"strings"
     9  
    10  	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
    11  	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/bloberror"
    12  	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blockblob"
    13  	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
    14  	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/service"
    15  	"github.com/Shopify/go-lua"
    16  	"github.com/treeverse/lakefs/pkg/block/azure"
    17  	"github.com/treeverse/lakefs/pkg/block/params"
    18  	"github.com/treeverse/lakefs/pkg/uri"
    19  )
    20  
    21  type Client struct {
    22  	ctx    context.Context
    23  	client *service.Client
    24  }
    25  
    26  func newBlobClient(ctx context.Context) lua.Function {
    27  	return func(l *lua.State) int {
    28  		storageAccount := lua.CheckString(l, 1)
    29  		accessKey := lua.CheckString(l, 2)
    30  		azClient, err := azure.BuildAzureServiceClient(params.Azure{
    31  			StorageAccount:   storageAccount,
    32  			StorageAccessKey: accessKey,
    33  		})
    34  		if err != nil {
    35  			panic(err)
    36  		}
    37  
    38  		client := &Client{
    39  			ctx:    ctx,
    40  			client: azClient,
    41  		}
    42  
    43  		l.NewTable()
    44  		functions := map[string]lua.Function{
    45  			"get_object":    client.GetObject,
    46  			"put_object":    client.PutObject,
    47  			"delete_object": client.DeleteObject,
    48  		}
    49  		for name, goFn := range functions {
    50  			l.PushGoFunction(goFn)
    51  			l.SetField(-2, name)
    52  		}
    53  		return 1
    54  	}
    55  }
    56  
    57  func (c *Client) GetObject(l *lua.State) int {
    58  	path := lua.CheckString(l, 1)
    59  	containerName, key := parsePath(l, path)
    60  	blobClient := c.getBlobClient(containerName, key)
    61  	downloadResponse, err := blobClient.DownloadStream(c.ctx, &azblob.DownloadStreamOptions{})
    62  
    63  	if err != nil {
    64  		if bloberror.HasCode(err, bloberror.BlobNotFound) {
    65  			l.PushString("")
    66  			l.PushBoolean(false) // exists
    67  			return 2
    68  		}
    69  		lua.Errorf(l, "%s", err.Error())
    70  		panic("unreachable")
    71  	}
    72  
    73  	data, err := io.ReadAll(downloadResponse.Body)
    74  	if err != nil {
    75  		lua.Errorf(l, "%s", err.Error())
    76  		panic("unreachable")
    77  	}
    78  	l.PushString(string(data))
    79  	l.PushBoolean(true) // exists
    80  	return 2
    81  }
    82  
    83  func (c *Client) PutObject(l *lua.State) int {
    84  	path := lua.CheckString(l, 1)
    85  	buf := strings.NewReader(lua.CheckString(l, 2))
    86  	containerName, key := parsePath(l, path)
    87  	blobClient := c.getBlobClient(containerName, key)
    88  	_, err := blobClient.UploadStream(c.ctx, buf, &azblob.UploadStreamOptions{})
    89  	if err != nil {
    90  		lua.Errorf(l, "azure client: (container %s) (key %s): %s", containerName, key, err.Error())
    91  		panic("unreachable")
    92  	}
    93  	return 0
    94  }
    95  
    96  func (c *Client) DeleteObject(l *lua.State) int {
    97  	path := lua.CheckString(l, 1)
    98  	containerName, key := parsePath(l, path)
    99  	blobClient := c.getBlobClient(containerName, key)
   100  	_, err := blobClient.Delete(c.ctx, nil)
   101  	if err != nil {
   102  		lua.Errorf(l, "%s", err.Error())
   103  		panic("unreachable")
   104  	}
   105  	return 0
   106  }
   107  
   108  func (c *Client) getBlobClient(container, blob string) *blockblob.Client {
   109  	return c.getContainerClient(container).NewBlockBlobClient(blob)
   110  }
   111  
   112  func (c *Client) getContainerClient(container string) *container.Client {
   113  	return c.client.NewContainerClient(container)
   114  }
   115  
   116  func parsePath(l *lua.State, path string) (string, string) {
   117  	// Extract containerName and key from path
   118  	containerName, key, found := strings.Cut(path, uri.PathSeparator)
   119  	if !found {
   120  		lua.Errorf(l, "azure client: invalid path, missing container name from path: %s", path)
   121  		panic("unreachable")
   122  	}
   123  	return containerName, key
   124  }
   125  
   126  func transformPathToAbfss(l *lua.State) int {
   127  	path := lua.CheckString(l, 1)
   128  	const numOfParts = 3
   129  	// Added adls for backwards compatibility in imports created pre fix of bug: https://github.com/treeverse/lakeFS/issues/7580
   130  	r := regexp.MustCompile(`^https://(\w+)\.(?:blob|adls)\.core\.windows\.net/([^/]*)/(.+)$`)
   131  	parts := r.FindStringSubmatch(path)
   132  	if len(parts) != numOfParts+1 {
   133  		lua.Errorf(l, "expected valid Azure https URL: %s", path)
   134  		panic("unreachable")
   135  	}
   136  	transformed := fmt.Sprintf("abfss://%s@%s.dfs.core.windows.net/%s", parts[2], parts[1], parts[3])
   137  	l.PushString(transformed)
   138  	return 1
   139  }