# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
import k1lib as k1, os
import k1lib.cli as cli
from collections import deque
from functools import lru_cache
from contextlib import contextmanager
__all__ = ["sql", "s3", "s3bucket"]
@contextmanager
def mysqlCnf(user, password, host, port):                                        # mysqlCnf
    fn = f"""[client]\nuser = "{user}"\npassword = "{password or ''}"\nhost = "{host}"\nport = "{port}" """ | cli.file() # mysqlCnf
    try: yield fn                                                                # mysqlCnf
    finally: os.remove(fn)                                                       # mysqlCnf
mysqlConn = k1.dep("mysql.connector", "mysql-connector-python", "https://pypi.org/project/mysql-connector-python/") # mysqlCnf
pgConn = k1.dep("psycopg2", "psycopg2-binary", "https://pypi.org/project/psycopg2/") # mysqlCnf
sqlite3 = k1.dep("sqlite3")                                                      # mysqlCnf
qD = {"my": "`", "pg": "", "lite": ""} # quote dict                              # mysqlCnf
class sql:                                                                       # sql
    def __init__(self, host="127.0.0.1", port=3306, user=None, password="", mode="my"): # sql
        """Creates a connection to a SQL database.
Example::

    s = sql("127.0.0.1") # creates a new sql object. Apparently, mysql and mysqldump on the command line bugs out if you put "localhost". If you put localhost here, it should work fine, but `sql(...) | toBytes()` and other functions that use external programs can deny you access
    s.refresh()          # refreshes connection any time you encounter strange bugs

    s | ls()                                                       # returns List[sqldb], lists out all databases
    s | toBytes()                                                  # returns Iterator[str], dumps every databases. Yes, it's "toBytes", but it has the feeling of serializing whatever the input is, so it's at least intuitive in that way
    "dump.sql" | s                                                 # restores the database using the dump file
    cat("dump.sql") | s                                            # restores the database using the dump file
    db1 = s["db1"]                                                 # returns sqldb, gets database named "db1"
    db1 | ls()                                                     # returns List[sqltable], list out all tables within this database
    db1 | toBytes()                                                # returns Iterator[str], dumps the database
    users = db1["user"]                                            # gets table named "user", short and simple

    db1.query("select * from users")                     # queries the database using your custom query
    db1.query("select * from users where user_id=%s", 3) # queries with prepared statement

    users.info()                                               # prints out the first 10 rows of the table and the table schema
    users.cols                                                 # returns table's columns as List[str]
    len(users)                                                 # returns number of rows
    users.query(...)                                           # can also do a custom query, just like with databases
    users[1]                                                   # grabs user with id 1. Returns None if user doesn't exist
    users[1:10]                                                # grabs users with id from 1 (inclusive) to 10 (exclusive)
    users[1].firstname = "Reimu"                               # sets the first name of user with id 1 to "Reimu"
    users[1] = {"firstname": "Reimu", "lastname": "Hakurei"}   # sets first name and last name of user with id 1
    users | toBytes()                                          # dumps this specific table, returns Iterator[str]
    users |  cat() | display()                                 # reads entire table, gets first 10 rows and displays it out
    users | (cat() | head(20)) | display()                     # reads first 20 rows only, then displays the first 10 rows. Query sent is "select * from user limit 20"
    users | (cat() | filt("x == 4", 3) | head(20)) | display() # grabs first 20 rows that has the 4th column equal to 4, then displays the first 10 rows. Query sent is "select user_id, address, balance, age from user where age = 4", assuming the table only has those columns

Philosophy for these methods is that they should be intuitive to look at and interact with, not for performance.
If performance is needed, just write raw queries or don't use a database and use applyCl() instead.

:param host: host name, ip address, or file name (in case of sqlite)
:param port: port at the host
:param uesr: database user name. If not specified then fallback to environment variable ``SQL_USER``, then ``USER``
:param password: database password. If not specified then assume database doesn't require one
:param mode: currently supports 3 values: "my" (MySQL), "pg" (PostgreSQL) and "lite" (SQLite)
"""                                                                              # sql
        self.host = host; self.port = port; self.user = user or os.environ.get("SQL_USER") or os.environ.get("USER") # sql
        self.password = password or os.environ.get("SQL_PASSWORD"); self.db = None; self.mode = mode; self.conn = None; self.refresh() # sql
    def refresh(self):                                                           # sql
        try: self.conn.close(); self.conn = None                                 # sql
        except: pass                                                             # sql
        kwargs = dict(host=self.host, port=self.port, user=self.user, password=self.password, database=self.db) # sql
        if self.mode == "my": self.conn = mysqlConn.connect(**kwargs, charset='utf8mb4', collation='utf8mb4_general_ci') # sql
        elif self.mode == "pg": self.conn = pgConn.connect(**kwargs)             # sql
        elif self.mode == "lite": self.conn = sqlite3.connect(self.host)         # sql
        else: raise Exception(f"Can only support 'my' (MySQL), 'pg' (PostgreSQL) and 'lite' (SQLite) for now, can't support {self.mode}") # sql
    def _changeDb(self, db):                                                     # sql
        if self.db != db:                                                        # sql
            if self.mode == "my": self.query(f"use `{db}`")                      # sql
            elif self.mode == "pg": self.db = db; self.refresh()                 # sql
            elif self.mode == "lite": pass                                       # sql
        self.db = db                                                             # sql
    def query(self, query, *args):                                               # sql
        cur = self.conn.cursor(); cur.execute(query, args)                       # sql
        try: ans = cur.fetchall()                                                # sql
        except: ans = None                                                       # sql
        cur.close(); self.conn.commit(); return ans                              # sql
    def _ls(self):                                                               # sql
        if self.mode == "my": return [sqldb(self, e[0]) for e in self.query("show databases")] # sql
        elif self.mode == "pg": return [sqldb(self, e[0]) for e in self.query("select datname from pg_database where datistemplate=false")] # sql
        elif self.mode == "lite": return [sqldb(self, "default")]                # sql
    def __repr__(self): return f"<sql mode={self.mode} host={self._host}>"       # sql
    @property                                                                    # sql
    def _host(self) -> str:                                                      # sql
        if self.mode == "lite": return self.host.split(os.sep)[-1]               # sql
        else: return f"{self.host}:{self.port}"                                  # sql
    def _cnfCtx(self): return mysqlCnf(self.user, self.password, self.host, self.port) # sql
    def _toBytes(self):                                                          # sql
        if self.mode == "my":                                                    # sql
            with self._cnfCtx() as fn: yield from None | cli.cmd(f"mysqldump --defaults-file={fn} --single-transaction --hex-blob --all-databases") # sql
        else: raise Exception(f"All databases dump of mode {self.mode} is not supported yet") # sql
    def __ror__(self, it): # restoring a backup                                  # sql
        if self.mode == "my":                                                    # sql
            def restore(fn):                                                     # sql
                with self._cnfCtx() as cnfFn: None | cli.cmd(f"mysql --defaults-file={cnfFn} < {fn}") | cli.ignore() # sql
            if isinstance(it, str):     restore(it)                              # sql
            else: fn = it | cli.file(); restore(fn); os.remove(fn)               # sql
        else: raise Exception(f"Restoring database from .sql file of mode {self.mode} is not supported yet") # sql
    def __len__(self): return len(self._ls())                                    # sql
    def __getitem__(self, idx):                                                  # sql
        if isinstance(idx, str):                                                 # sql
            lsres = [x for x in self._ls() if x.name == idx]                     # sql
            if len(lsres) == 1: return lsres[0]                                  # sql
            else: raise Exception(f"Database '{idx}' does not exist, can't retrieve it") # sql
        return self._ls()[idx]                                                   # sql
    def __iter__(self): return iter(self._ls())                                  # sql
    def __delitem__(self, idx):                                                  # sql
        if isinstance(idx, str):                                                 # sql
            lsres = [x for x in self._ls() if x.name == idx]                     # sql
            if len(lsres) == 1:                                                  # sql
                if self.mode == "my" or self.mode == "pg": return self.query(f"drop database {idx}") # sql
                else: raise Exception("Currently only support dropping databases in mysql and postgresql") # sql
            else: raise Exception(f"Database '{idx}' does not exist, can't delete it") # sql
        else: raise Exception("Only supports deleting with table names. Use this like `del db['some_table']`") # sql
class sqldb:                                                                     # sqldb
    def __init__(self, sql:sql, name:str):                                       # sqldb
        """A sql database representation. Not expected to be instatiated by you. See also: :class:`sql`""" # sqldb
        self.sql = sql; self.name = name                                         # sqldb
    def query(self, query, *args): self.sql._changeDb(self.name); return self.sql.query(query, *args) # sqldb
    def _ls(self):                                                               # sqldb
        if self.sql.mode == "my": return [sqltable(self.sql, self, e[0]) for e in self.query(f"show tables")] # sqldb
        if self.sql.mode == "pg": return [sqltable(self.sql, self, e[0]) for e in self.query(f"select table_name from information_schema.tables")] # sqldb
        if self.sql.mode == "lite": return [sqltable(self.sql, self, e[0]) for e in self.query("select name from sqlite_master where type='table'")] # sqldb
    def __repr__(self): return f"<sqldb host={self.sql._host} db={self.name}>"   # sqldb
    def _toBytes(self):                                                          # sqldb
        if self.sql.mode == "my":                                                # sqldb
            with self.sql._cnfCtx() as fn: yield from None | cli.cmd(f"mysqldump --defaults-file={fn} --single-transaction --hex-blob --databases {self.name}") # sqldb
        else: raise Exception(f"Database dump of mode {self.sql.mode} is not supported yet") # sqldb
    def __ror__(self, it): return self.sql.__ror__(it)                           # sqldb
    def __len__(self): return len(self._ls())                                    # sqldb
    def __getitem__(self, idx):                                                  # sqldb
        if isinstance(idx, str):                                                 # sqldb
            lsres = [x for x in self._ls() if x.name == idx]                     # sqldb
            if len(lsres) == 1: return lsres[0]                                  # sqldb
            else: raise Exception(f"Table '{idx}' does not exist, can't retrieve it") # sqldb
        return self._ls()[idx]                                                   # sqldb
    def __iter__(self): return iter(self._ls())                                  # sqldb
    def __delitem__(self, idx):                                                  # sqldb
        if isinstance(idx, str):                                                 # sqldb
            lsres = [x for x in self._ls() if x.name == idx]                     # sqldb
            if len(lsres) == 1: return self.query(f"drop table {idx}")           # sqldb
            else: raise Exception(f"Table '{idx}' does not exist, can't delete it") # sqldb
        else: raise Exception("Only supports deleting with table names. Use this like `del db['some_table']`") # sqldb
class sqltable:                                                                  # sqltable
    def __init__(self, sql, sqldb, name:str):                                    # sqltable
        """A sql table representation. Not expected to be instantiated by you. See also: :class:`sql`""" # sqltable
        self.sql = sql; self.sqldb = sqldb; self.name = name; self._cols = None  # sqltable
    def _cat(self, ser):                                                         # sqltable
        cols = self.cols; _2 = [] # clis that can't be optimized, stashed away to be merged with ser later on # sqltable
        q = qD[self.sql.mode]; clis = deque(ser.clis)                            # sqltable
        o1 = None # cut() opt                                                    # sqltable
        o2 = None # head() opt                                                   # sqltable
        o3 = [] # filt() opt                                                     # sqltable
        while len(clis) > 0:                                                     # sqltable
            c = clis.popleft()                                                   # sqltable
            if isinstance(c, cli.filt): _2.append(c); break # TODO: add optimizations for filt # sqltable
            elif o2 is None and isinstance(c, cli.head):                         # sqltable
                if round(c.n) != c.n or c.n < 0 or c.inverted or c.n == None: _2.append(c); break # sqltable
                else: o2 = f"limit {c.n}"; continue                              # sqltable
            elif o1 is None and isinstance(c, cli.cut):                          # sqltable
                if isinstance(c.columns, slice): _2.append(c); o1 = 0; continue  # sqltable
                else:                                                            # sqltable
                    o1 = ", ".join([f"{q}{c}{q}" for c in cols | cli.rows(*c.columns)]) # sqltable
                    if len(c.columns) == 1: _2.append(cli.item().all() | cli.aS(list)) # sqltable
            else: _2.append(c); break                                            # sqltable
        o1 = o1 or ", ".join([f"{q}{c}{q}" for c in cols])                       # sqltable
        query = f"select {o1} from {q}{self.name}{q} {o2 or ''}"#; print(f"query: {query}"); return [] # sqltable
        sql = self.sql; return [sqlrow(sql, self, row) for row in self.sqldb.query(query) | cli.serial(*_2, *clis)] # sqltable
    @property                                                                    # sqltable
    def cols(self):                                                              # sqltable
        """Get column names"""                                                   # sqltable
        if not self._cols: self._cols = self.describe()[1:] | cli.cut({"my": 0, "pg": 0, "lite": 1}[self.sql.mode]) | cli.deref() # sqltable
        return self._cols                                                        # sqltable
    @lru_cache                                                                   # sqltable
    def describe(self):                                                          # sqltable
        if self.sql.mode == "my": return self.sqldb.query(f"describe `{self.name}`") | cli.insert(["Field", "Type", "Null", "Key", "Default", "Extra"]) | cli.deref() # sqltable
        if self.sql.mode == "pg": return self.sqldb.query(f"select column_name, data_type, is_nullable, column_default, ordinal_position from information_schema.columns where table_name='{self.name}'") | cli.insert(["column_name", "data_type", "is_nullable", "column_default", "ordinal_position"]) | cli.deref() # sqltable
        if self.sql.mode == "lite": return self.sqldb.query(f"pragma table_info([{self.name}])") | cli.insert(["cid", "name", "type", "notnull", "dflt_value", "pk"]) | cli.deref() # sqltable
    def insert(self, **kwargs):                                                  # sqltable
        """Inserts a row.
Example::

    table = ...
    table.insert(firstname="Yuyuko", lastname="Saigyouji")
"""                                                                              # sqltable
        keys = ", ".join([f"`{x}`" for x in kwargs.keys()])                      # sqltable
        values = ", ".join(['%s']*len(kwargs))                                   # sqltable
        self.query(f"insert into {self.name} ({keys}) values ({values})", *kwargs.values()) # sqltable
    def _update(self, idx, **kwargs):                                            # sqltable
        """Updates a row.
Example::

    rable = ...
    table.update(3, firstname="Youmu", lastname="Konpaku")
"""                                                                              # sqltable
        p1 = ", ".join([f"`{k}` = %s" for k in kwargs.keys()])                   # sqltable
        self.query(f"update {self.name} set {p1} where {self.cols[0]} = {idx}", *kwargs.values()) # sqltable
    def info(self, out=False):                                                   # sqltable
        """Preview table.
Example::

    table = ...
    table.info()

:param out: if True, returns a list of lines instead of printing them out"""     # sqltable
        def gen():                                                               # sqltable
            print(f"Table `{self.name}`\n")                                      # sqltable
            desc = self.describe() | cli.deref(); cols = self.cols; q = qD[self.sql.mode]; s = ", ".join([f"{q}{e}{q}" for e in cols]) # sqltable
            self.sqldb.query(f"select {s} from {q}{self.name}{q} limit 9") | (cli.aS(repr) | cli.head(50)).all(2) | cli.insert(cols) | cli.display(); print("") # sqltable
            desc | cli.display(None)                                             # sqltable
        if out:                                                                  # sqltable
            with k1.captureStdout() as out: gen()                                # sqltable
            return out()                                                         # sqltable
        else: gen()                                                              # sqltable
    def query(self, query, *args): return self.sqldb.query(query, *args)         # sqltable
    def __repr__(self): return f"<sqltable host={self.sql._host} db={self.sqldb.name} table={self.name}>" # sqltable
    def _toBytes(self):                                                          # sqltable
        if self.sql.mode == "my":                                                # sqltable
            with self.sql._cnfCtx() as fn: yield from None | cli.cmd(f"mysqldump --defaults-file={fn} --single-transaction --hex-blob {self.sqldb.name} {self.name}") # sqltable
        else: raise Exception(f"Table dump of mode {self.sql.mode} is not supported yet") # sqltable
    def __ror__(self, it): return self.sql.__ror__(it)                           # sqltable
    def __len__(self): return self.query(f"select count(*) from {self.name}")[0][0] # sqltable
    def __getitem__(self, idx):                                                  # sqltable
        idName = self.cols[0]                                                    # sqltable
        if isinstance(idx, int):                                                 # sqltable
            res = self.query(f"select * from {self.name} where {idName} = {idx}") # sqltable
            if len(res) == 0: return None                                        # sqltable
            return sqlrow(self.sql, self, res[0])                                # sqltable
        elif isinstance(idx, slice):                                             # sqltable
            idxs = list(range(idx.stop if idx.stop is not None else (self.query(f"select max({idName}) from {self.name}")[0][0]+1))[idx]) # sqltable
            res = self.query(f"select * from {self.name} where {idName} in ({', '.join([str(x) for x in idxs])})") # sqltable
            d = {row[0]:row for row in res}; cols = self.cols; return [sqlrow(self.sql, self, d.get(idx, None)) for idx in idxs] # sqltable
        else: raise Exception("Only support table indexing of integers or slices") # sqltable
    def __setitem__(self, idx, value):                                           # sqltable
        if not isinstance(value, dict): raise Exception("Only accepts setting elements with dicts") # sqltable
        if self[idx] is None: raise Exception(f"Can't set element with id {idx}, it doesn't seem to exist. Use table.insert(col1=value1, col2=value2) instead") # sqltable
        self._update(idx, **value)                                               # sqltable
    def __iter__(self): return iter(self | cli.cat())                            # sqltable
    def __delitem__(self, idx):                                                  # sqltable
        idName = self.cols[0]                                                    # sqltable
        if isinstance(idx, int):                                                 # sqltable
            res = self[idx]                                                      # sqltable
            if res: self.query(f"delete from {self.name} where {idName} = %s", res[0]) # sqltable
        elif isinstance(idx, slice):                                             # sqltable
            res = self[idx]; ids = ", ".join([row[0] for row in res if row])     # sqltable
            self.query(f"delete from {self.name} where {idName} in {ids}")       # sqltable
        else: raise Exception("Only support table indexing of integers or slices") # sqltable
class sqlrow:                                                                    # sqlrow
    def __init__(self, sql, sqltable, row):                                      # sqlrow
        self._ab_sentinel = True                                                 # sqlrow
        self.sql = sql; self.sqltable = sqltable                                 # sqlrow
        self.row = row; self._d = None                                           # sqlrow
        self._ab_sentinel = False                                                # sqlrow
    def __getitem__(self, idx): return self.row[idx]                             # sqlrow
    def __getattr__(self, attr):                                                 # sqlrow
        if self._d is None: self._d = {c:r for c,r in zip(self.sqltable.cols, self.row)} # sqlrow
        return self._d[attr]                                                     # sqlrow
    def __setattr__(self, attr, value):                                          # sqlrow
        if attr == "_ab_sentinel": self.__dict__[attr] = value                   # sqlrow
        else:                                                                    # sqlrow
            if self._ab_sentinel or attr in self.__dict__: self.__dict__[attr] = value # sqlrow
            else: return self.sqltable._update(self.row[0], **{attr: value})     # sqlrow
    def __len__(self): return len(self.row)                                      # sqlrow
    def __repr__(self): return f"{self.row}"                                     # sqlrow
boto3 = k1.dep("boto3")                                                          # sqlrow
class s3:                                                                        # s3
    def __init__(self, client):                                                  # s3
        """Represents an S3 client.
Example::

    client = boto3.client("s3", ...)            # put your credentials and details here
    db = s3(client)                             # creates an S3 manager
    db | ls()                                   # lists all buckets accessible

    bucket = db | ls() | item()                 # grabs the first bucket, returns object of type s3bucket
    bucket = s3bucket(client, "bucket-name")    # or you can instantiate the bucket directly
    bucket | ls()                               # lists all objects within this bucket
    bucket | ls() | grep("\\.so")               # grabs all .so files from the bucket

    obj = bucket | ls() | item()                # grabs the first object within this bucket, returns object of type s3obj
    obj.key, obj.size, obj.lastModified         # some fields directly accessible

This mostly offers interoperability with ls() and cat(), so that you can
write relatively intuitive code, but fundamentally provides no upsides"""        # s3
        self.client = client                                                     # s3
    def _ls(self): return [s3bucket(self.client, x["Name"]) for x in self.client.list_buckets()["Buckets"]] # s3
    def __repr__(self): return f"<kaws.s3 client>"                               # s3
class s3bucket:                                                                  # s3bucket
    def __init__(self, client, name:str):                                        # s3bucket
        """Represents an S3 bucket.
See also: :class:`s3`"""                                                         # s3bucket
        self.client = client; self.name = name                                   # s3bucket
    def _ls(self): client = self.client; name = self.name; return [s3obj(client, name, data) for data in self.client.list_objects(Bucket=name).get("Contents", [])] # s3bucket
    def __repr__(self): return f"<s3bucket name='{self.name}'>"                  # s3bucket
class s3obj:                                                                     # s3obj
    def __init__(self, client, bucket:str, data):                                # s3obj
        """Represents an S3 object. Not intended to be instantiated directly.
See also: :class:`s3`"""                                                         # s3obj
        self.client = client; self.bucket = bucket                               # s3obj
        self.key  = data["Key"];  self.lastModified = data["LastModified"]       # s3obj
        self.size = data["Size"]; self.storageClass = data["StorageClass"]       # s3obj
    def __repr__(self): return f"<s3obj bucket='{self.bucket}' key='{self.key}' size={fmt.size(self.size)}>" # s3obj
    def _cat(self, kwargs):                                                      # s3obj
        if kwargs["text"] is not False: raise Exception(f"s3obj does not support `cat(text=True)`. Please use `cat(text=False)` instead") # s3obj
        if kwargs["chunks"]: raise Exception(f"s3obj does not support `cat(chunks=True)`") # s3obj
        sB = kwargs["sB"]; eB = kwargs["eB"]                                     # s3obj
        if eB < 0: eB = self.size                                                # s3obj
        return self.client.get_object(Bucket=self.bucket, Key=self.key, Range=f'bytes={sB}-{eB-1}')["Body"].read() # s3obj
