start working on making this usable

Author: withinboredom

atlas/authorization.go

@@ -21,14 +21,16 @@ package atlas
 import (
 	"go.uber.org/zap"
 	"slices"
+	"strings"
 	"sync"
 	"time"
 	"zombiezen.com/go/sqlite"
 )
 
 type Authorizer struct {
-	boundTime time.Duration
-	mu        sync.RWMutex
+	boundTime  time.Duration
+	mu         sync.RWMutex
+	LastTables map[string]struct{}
 }
 
 var writeOps []sqlite.OpType = []sqlite.OpType{
@@ -74,6 +76,7 @@ func (a *Authorizer) Reset() {
 	defer a.mu.Unlock()
 
 	a.boundTime = 0
+	a.LastTables = make(map[string]struct{})
 }
 
 func (a *Authorizer) isJournalModeChange(action sqlite.Action) bool {
@@ -91,6 +94,11 @@ func (a *Authorizer) isAtlasChange(action sqlite.Action) bool {
 func (a *Authorizer) Authorize(action sqlite.Action) sqlite.AuthResult {
 	Logger.Info("Auth", zap.Any("action", action), zap.String("table", action.Table()))
 
+	if action.Table() != "" && action.Database() != "" && action.Table() != "sqlite_master" {
+		name := strings.ToUpper(action.Database()) + "." + strings.ToUpper(action.Table())
+		a.LastTables[name] = struct{}{}
+	}
+
 	if a.isJournalModeChange(action) || a.isAtlasChange(action) {
 		return sqlite.AuthResultDeny
 	}

atlas/consensus/server.go

@@ -388,7 +388,7 @@ func (s *Server) applyMigration(migrations []*Migration, commitConn *sqlite.Conn
 	return nil
 }
 
-func constructCurrentNode() *Node {
+func ConstructCurrentNode() *Node {
 	return &Node{
 		Id:      atlas.CurrentOptions.ServerId,
 		Address: atlas.CurrentOptions.AdvertiseAddress,
@@ -515,7 +515,7 @@ VALUES (:id, :address, :port, :region, 1, current_timestamp, 0)`, conn, false, a
 	}
 
 	mreq := &WriteMigrationRequest{
-		Sender:    constructCurrentNode(),
+		Sender:    ConstructCurrentNode(),
 		Migration: migration,
 	}
 
@@ -635,7 +635,7 @@ func SendGossip(ctx context.Context, req *GossipMigration, conn *sqlite.Conn) er
 		Table:             req.GetTable(),
 		PreviousMigration: req.GetPreviousMigration(),
 		Ttl:               req.GetTtl() - 1,
-		Sender:            constructCurrentNode(),
+		Sender:            ConstructCurrentNode(),
 	}
 
 	wg := sync.WaitGroup{}

atlas/pool.go

@@ -23,7 +23,6 @@ import (
 	_ "embed"
 	"runtime"
 	"strings"
-	"sync"
 	"zombiezen.com/go/sqlite"
 	"zombiezen.com/go/sqlite/sqlitemigration"
 )
@@ -34,27 +33,18 @@ var migrations string
 var Pool *sqlitemigration.Pool
 var MigrationsPool *sqlitemigration.Pool
 
-var authorizers = map[*sqlite.Conn]*Authorizer{}
-
 // CreatePool creates a new connection pool for the database and the migrations database.
 func CreatePool(options *Options) {
 	if Pool != nil {
 		return
 	}
 
-	authorizers = make(map[*sqlite.Conn]*Authorizer)
-	amu := sync.Mutex{}
-
 	Pool = sqlitemigration.NewPool(options.DbFilename, sqlitemigration.Schema{}, sqlitemigration.Options{
 		Flags:    sqlite.OpenReadWrite | sqlite.OpenCreate | sqlite.OpenWAL,
 		PoolSize: runtime.NumCPU() * 2,
 		PrepareConn: func(conn *sqlite.Conn) (err error) {
 			auth := &Authorizer{}
 
-			amu.Lock()
-			authorizers[conn] = auth
-			amu.Unlock()
-
 			err = conn.SetAuthorizer(auth)
 			if err != nil {
 				return

atlas/socket/handler.go

@@ -48,6 +48,7 @@ type Socket struct {
 	streams       []*sqlite.Stmt
 	principals    []*consensus.Principal
 	timeout       time.Duration
+	authorizer    *atlas.Authorizer
 }
 
 func (s *Socket) Cleanup() {
@@ -160,6 +161,9 @@ func (s *Socket) rollbackAutoTransaction(ctx context.Context, err error) error {
 }
 
 func (s *Socket) setTimeout(t time.Duration) error {
+	if t == 0 {
+		return s.conn.SetDeadline(time.Time{})
+	}
 	return s.conn.SetDeadline(time.Now().Add(t))
 }
 
@@ -247,10 +251,21 @@ ready:
 	}
 	defer atlas.Pool.Put(s.sql)
 
+	s.authorizer = &atlas.Authorizer{}
+
+	err = s.sql.SetAuthorizer(s.authorizer)
+	if err != nil {
+		atlas.Logger.Error("Error setting authorizer", zap.Error(err))
+		return
+	}
+	defer s.sql.SetAuthorizer(nil)
+
 	if err = s.setTimeout(s.timeout); err != nil {
 		return
 	}
 
+	var changes []*consensus.Migration
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -283,7 +298,30 @@ ready:
 				_, err = s.PerformExecute(ctx, cmd)
 				goto handleError
 			case "QUERY":
-				_, err = s.PerformQuery(ctx, cmd)
+				var q *Query
+				q, err = s.PerformQuery(ctx, cmd)
+				if err != nil {
+					goto handleError
+				}
+				if len(q.tables) > 0 {
+					// we need ownership over all tables
+					for _, table := range q.tables {
+						qm := consensus.GetDefaultQuorumManager(ctx)
+						var qu consensus.Quorum
+						qu, err = qm.GetQuorum(ctx, table.GetName())
+						if err != nil {
+							err = makeFatal(err)
+							goto handleError
+						}
+
+						result, err = qu.StealTableOwnership(ctx, &consensus.StealTableOwnershipRequest{
+							Sender: consensus.ConstructCurrentNode(),
+							Reason: consensus.StealReason_schemaReason,
+							Table:  nil,
+						})
+					}
+				}
+
 				goto handleError
 			case "FINALIZE":
 				err = s.PerformFinalize(cmd)

atlas/socket/query.go

@@ -19,33 +19,55 @@
 package socket
 
 import (
+	"github.com/bottledcode/atlas-db/atlas"
 	"github.com/bottledcode/atlas-db/atlas/commands"
+	"github.com/bottledcode/atlas-db/atlas/consensus"
+	"github.com/bottledcode/atlas-db/atlas/operations"
+	"go.uber.org/zap"
 	"zombiezen.com/go/sqlite"
 )
 
 type Query struct {
-	stmt  *sqlite.Stmt
-	query *commands.SqlCommand
+	stmt   *sqlite.Stmt
+	query  *commands.SqlCommand
+	tables []*consensus.Table
 }
 
-func ParseQuery(cmd *commands.CommandString) (*Query, error) {
+func ParseQuery(cmd *commands.CommandString) (query *Query, err error) {
 	if err := cmd.CheckMinLen(2); err != nil {
 		return nil, err
 	}
 
 	q := cmd.From(1)
 
-	return &Query{
+	query = &Query{
 		query: q,
 		stmt:  nil,
-	}, nil
+	}
+
+	return query, nil
 }
 
 func (q *Query) Handle(s *Socket) (err error) {
+	s.authorizer.Reset()
 	q.stmt, _, err = s.sql.PrepareTransient(q.query.Raw())
 	if err != nil {
 		return makeFatal(err)
 	}
+	tables := s.authorizer.LastTables
+	atlas.Logger.Info("tables", zap.Any("tables", tables))
+
+	if !q.query.IsQueryReadOnly() {
+		first, _ := q.query.SelectNormalizedCommand(0)
+		if first == "CREATE" {
+			q.tables, err = operations.CreateTable(q.query)
+			if err != nil {
+				return nil
+			}
+		}
+
+		// todo: handle alters
+	}
 
 	s.streams = append(s.streams, q.stmt)
 	return nil

docs/reference/commands.md

@@ -122,59 +122,127 @@ rules.
    ```
    Prepares a query for execution, associating it with an identifier.
 
+   Returns: OK or ERROR
+
 2. **Execute a Prepared Query:**
    ```
    EXECUTE [ID]\r\n
    ```
    Executes a previously prepared query.
 
+   Returns: Stream + Metadata + OK
+
 3. **Perform a One-Off Query:**
    ```
    QUERY [RAW QUERY]\r\n
    ```
    Executes a single query without preparation.
 
+   Returns: Stream + Metadata + OK
+
 4. **Remove a Prepared Query:**
    ```
    FINALIZE [ID]\r\n
    ```
    Removes a previously prepared query from memory.
 
+   Returns: OK or ERROR
+
 5. **Bind a Value to a Prepared Query:**
    ```
    BIND [ID] [PARAM] [TYPE] [VALUE]\r\n
    ```
    Binds a value to a parameter in a prepared query.
 
-6. **Transaction Commands:**
+   Returns: OK or ERROR
+
+6. **Begin:**
+
+   ```
+   BEGIN\r\n
+   ```
+   Starts a readonly transaction.
+
+   Returns: OK or ERROR
+
+7. **BEGIN IMMEDIATE**
+   ```
+   BEGIN IMMEDIATE WITH [TABLE (table_name, other table)] [VIEW (view_name, other_view)] [TRIGGER (trigger_name)]\r\n
+   ```
+   Starts a writable transaction with write access to the declared objects.
+   Note that if any tables require a principal, the principal must be declared before starting the transaction.
+
+   Returns: OK or ERROR
+
+8. **COMMIT**
    ```
-   BEGIN IMMEDIATE\r\n
    COMMIT\r\n
+   ```
+   Commits the current transaction.
+
+   Returns: OK or ERROR
+
+9. **ROLLBACK**
+   ```
    ROLLBACK\r\n
+   ```
+   Rollback the current transaction.
+
+   Returns: OK or ERROR
+
+10. **SAVEPOINT**
+   ```
    SAVEPOINT [name]\r\n
-   RELEASE [name]\r\n
    ```
-   Manages database transactions, including savepoints for partial rollbacks.
+   Create a savepoint in the current transaction.
+
+   Returns: OK or ERROR
 
-7. **Pragmas:**
+11. **RELEASE**
    ```
-   PRAGMA [name]=[VALUE]\r\n
-   PRAGMA [name]\r\n
+   RELEASE [name]\r\n
    ```
-   Configures or queries node-specific settings.
+   Release a savepoint in the current transaction.
+
+12. **Pragmas:**
+    ```
+    PRAGMA [name]=[VALUE]\r\n
+    PRAGMA [name]\r\n
+    ```
+    Configures or queries node-specific settings.
+
+13. **Principle Command:**
+    ```
+    PRINCIPLE [principle_name] [id]\r\n
+    ```
+    Sets the principle context for row-level security.
+    A valid principle must be established before interacting with tables that use row-level security.
+    It cannot be used in a transaction.
+
+    Returns: OK or ERROR
 
-8. **Principle Command:**
+14. **Scroll Command:**
+    ```
+    SCROLL [StreamID] [Count]\r\n
+    ```
+    Fetches a specific number of rows from a stream identified by `StreamID`.
+
+    RETURNS: ROWS + METADATA + OK
+15. **RESET:**
    ```
-   PRINCIPLE [principle_name] [id]\r\n
+   RESET [ID]\r\n
    ```
-   Sets the principle context for row-level security. A valid principle must be established before interacting with
-   tables that use row-level security.
+   Resets a prepared query, so it can be run again.
+
+   Returns: OK or ERROR
 
-9. **Scroll Command:**
+16. **CLEAR BINDINGS:**
    ```
-   SCROLL [StreamID] [Count]\r\n
+   CLEARBINDINGS [ID]\r\n
    ```
-   Fetches a specific number of rows from a stream identified by `StreamID`.
+   Clears the bindings for a prepared query.
+
+   Returns: OK or ERROR
 
 ---